<commit_before>import pandas as pd
import utils
import json
import pint


class DB(object):
	def __init__(self):
		self._groups = utils.SafeDict()
		self.planes = utils.AnnotatedDataFrame({
			'run': ('Run', None),
			'frame': ('Frame number', None),
			'name': ('Plane name', None),
			'normal_x': ('Normal vector: x component', None),
			'normal_y': ('Normal vector: y component', None),
			'normal_z': ('Normal vector: z component', None),
			'support_x': ('Support point: x component', 'angstrom'),
			'support_y': ('Support point: y component', 'angstrom'),
			'support_z': ('Support point: z component', 'angstrom'),
		})
		self.distances = utils.AnnotatedDataFrame({
			'run': ('Run', None),
			'frame': ('Frame number', None),
			'name': ('Plane name', None),
			'atom1': ('First atom index', None),
			'atom2': ('Second atom index', None),
			'dist': ('Distance', 'angstrom')
		})

	@property
	def groups(self):
		""" Known static atom groups.
		:return: Dictionary with group names as key, list of 0-based atom indices as value.
		"""
		return self._groups

	def write(self, fh):
		""" Writes the database to disk or stream.
		:param fh: File handle or filename.
		:return:
		"""
		if not hasattr(fh, 'write'):
			fh = open(fh, 'w')

		data = dict()

		# groups
		data['groups'] = self.groups
		# planes
		data['planes'] = self.planes.to_dict()
		data['planes-meta'] = self.planes.annotations_to_dict()
		# distances
		data['distances'] = self.distances.to_dict()
		data['distances-meta'] = self.distances.annotations_to_dict()

		# finalise
		fh.write(json.dumps(data, separators=(',', ':')))

	def read(self, fh):
		""" Reads the database from disk or stream.
		:param fh: File handle or filename.
		:return:
		"""
		if not hasattr(fh, 'write'):
			fh = open(fh, 'w')

		data = json.load(fh)

		# groups
		self._groups = utils.SafeDict(data['groups'])
		# planes
		self.planes = utils.AnnotatedDataFrame(data['planes-meta'], data['planes'])
		# distances
		self.distances = utils.AnnotatedDataFrame(data['distances-meta'], data['distances'])

		# cleanup
		fh.close()


class DatabaseProvider(object):
	@staticmethod
	def get_annotated_collections():
		return 'atommeta', 'atomframemeta', 'framemeta'


class MemoryDatabaseProvider(DatabaseProvider):
	def __init__(self, run, atom, atommeta, atomframe, atomframemeta, frame, framemeta):
		self.run = Map(run)
		self.atom = pd.DataFrame(atom)
		self.atommeta = Map(atommeta)
		self.atomframe = pd.DataFrame(atomframe)<commit_msg>Fixed file read open mode.<commit_after>import pandas as pd
import utils
import json
import pint


class DB(object):
	def __init__(self):
		self._groups = utils.SafeDict()
		self.planes = utils.AnnotatedDataFrame({
			'run': ('Run', None),
			'frame': ('Frame number', None),
			'name': ('Plane name', None),
			'normal_x': ('Normal vector: x component', None),
			'normal_y': ('Normal vector: y component', None),
			'normal_z': ('Normal vector: z component', None),
			'support_x': ('Support point: x component', 'angstrom'),
			'support_y': ('Support point: y component', 'angstrom'),
			'support_z': ('Support point: z component', 'angstrom'),
		})
		self.distances = utils.AnnotatedDataFrame({
			'run': ('Run', None),
			'frame': ('Frame number', None),
			'name': ('Plane name', None),
			'atom1': ('First atom index', None),
			'atom2': ('Second atom index', None),
			'dist': ('Distance', 'angstrom')
		})

	@property
	def groups(self):
		""" Known static atom groups.
		:return: Dictionary with group names as key, list of 0-based atom indices as value.
		"""
		return self._groups

	def write(self, fh):
		""" Writes the database to disk or stream.
		:param fh: File handle or filename.
		:return:
		"""
		if not hasattr(fh, 'write'):
			fh = open(fh, 'w')

		data = dict()

		# groups
		data['groups'] = self.groups
		# planes
		data['planes'] = self.planes.to_dict()
		data['planes-meta'] = self.planes.annotations_to_dict()
		# distances
		data['distances'] = self.distances.to_dict()
		data['distances-meta'] = self.distances.annotations_to_dict()

		# finalise
		fh.write(json.dumps(data, separators=(',', ':')))

	def read(self, fh):
		""" Reads the database from disk or stream.
		:param fh: File handle or filename.
		:return:
		"""
		if not hasattr(fh, 'read'):
			fh = open(fh)

		data = json.load(fh)

		# groups
		self._groups = utils.SafeDict(data['groups'])
		# planes
		self.planes = utils.AnnotatedDataFrame(data['planes-meta'], data['planes'])
		# distances
		self.distances = utils.AnnotatedDataFrame(data['distances-meta'], data['distances'])

		# cleanup
		fh.close()


class DatabaseProvider(object):
	@staticmethod
	def get_annotated_collections():
		return 'atommeta', 'atomframemeta', 'framemeta'


class MemoryDatabaseProvider(DatabaseProvider):
	def __init__(self, run, atom, atommeta, atomframe, atomframemeta, frame, framemeta):
		self.run = Map(run)
		self.atom = pd.DataFrame(atom)
		self.atommeta = Map(atommeta)
		self.atomframe = pd.DataFrame(atomframe)
{'fraction_non_alphanumeric': 0.1652295969520754, 'fraction_numerical': 0.00120312813314618, 'mean_word_length': 4.357679914070892, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6411071', 'n_tokens_mistral': 1882, 'n_tokens_neox': 1672, 'n_words': 448}