Merge pull request #12 from samblau/BonDNet

BonDNet

HiPRGen/lmdb_dataset.py

@@ -0,0 +1,329 @@
+#give dgl graphs, reaction features, meta. write them into lmdb file.
+#1. check expend lmdb reasonably
+
+#give dgl graphs, reaction features, meta. write them into lmdb file.
+#1. check expend lmdb reasonably
+
+from torch.utils.data import Dataset
+from pathlib import Path
+import numpy as np
+import pickle
+import lmdb
+from torch.utils.data import random_split
+import multiprocessing as mp
+import os
+import pickle
+from tqdm import tqdm
+import glob 
+
+
+class LmdbDataset(Dataset):
+ """
+ Dataset class to 
+ 1. write Reaction networks objecs to lmdb
+ 2. load lmdb files
+ """
+ def __init__(self, config, transform=None):
+ super(LmdbDataset, self).__init__()
+
+ self.config = config
+ self.path = Path(self.config["src"])
+
+ #Get metadata in case
+ #self.metadata_path = self.path.parent / "metadata.npz"
+ self.env = self.connect_db(self.path)
+
+ # If "length" encoded as ascii is present, use that
+ # If there are additional properties, there must be length.
+ length_entry = self.env.begin().get("length".encode("ascii"))
+ if length_entry is not None:
+ num_entries = pickle.loads(length_entry)
+ else:
+ # Get the number of stores data from the number of entries
+ # in the LMDB
+ num_entries = self.env.stat()["entries"]
+
+ self._keys = list(range(num_entries))
+ self.num_samples = num_entries
+
+ #Get portion of total dataset
+ self.sharded = False
+ if "shard" in self.config and "total_shards" in self.config:
+ self.sharded = True
+ self.indices = range(self.num_samples)
+ # split all available indices into 'total_shards' bins
+ self.shards = np.array_split(
+ self.indices, self.config.get("total_shards", 1)
+ )
+ # limit each process to see a subset of data based off defined shard
+ self.available_indices = self.shards[self.config.get("shard", 0)]
+ self.num_samples = len(self.available_indices)
+
+ #TODO
+ self.transform = transform
+
+ def __len__(self):
+ return self.num_samples
+
+ def __getitem__(self, idx):
+ # if sharding, remap idx to appropriate idx of the sharded set
+ if self.sharded:
+ idx = self.available_indices[idx]
+
+ #!CHECK, _keys should be less then total numbers of keys as there are more properties.
+ datapoint_pickled = self.env.begin().get(
+ f"{self._keys[idx]}".encode("ascii")
+ )
+
+ data_object = pickle.loads(datapoint_pickled)
+
+ #TODO
+ if self.transform is not None:
+ data_object = self.transform(data_object)
+
+ return data_object
+
+ def connect_db(self, lmdb_path=None):
+ env = lmdb.open(
+ str(lmdb_path),
+ subdir=False,
+ readonly=True,
+ lock=False,
+ readahead=True,
+ meminit=False,
+ max_readers=1,
+ )
+ return env
+
+ def close_db(self):
+ if not self.path.is_file():
+ for env in self.envs:
+ env.close()
+ else:
+ self.env.close()
+
+ def get_metadata(self, num_samples=100):
+ pass
+
+ @property
+ def dtype(self):
+ dtype = self.env.begin().get("dtype".encode("ascii"))
+ return pickle.loads(dtype)
+
+ @property
+ def feature_size(self):
+ feature_size = self.env.begin().get("feature_size".encode("ascii"))
+ return pickle.loads(feature_size)
+
+ @property
+ def feature_name(self):
+ feature_name = self.env.begin().get("feature_name".encode("ascii"))
+ return pickle.loads(feature_name)
+
+ @property
+ def mean(self):
+ mean = self.env.begin().get("mean".encode("ascii"))
+ return pickle.loads(mean)
+
+ @property
+ def std(self):
+ std = self.env.begin().get("std".encode("ascii"))
+ return pickle.loads(std)
+
+
+def divide_to_list(a, b):
+ quotient = a // b
+ remainder = a % b
+
+ result = []
+ for i in range(b):
+ increment = 1 if i < remainder else 0
+ result.append(quotient + increment)
+
+ return result
+
+def cleanup_lmdb_files(directory, pattern):
+ """
+ Cleans up files matching the given pattern in the specified directory.
+ """
+ file_list = glob.glob(os.path.join(directory, pattern))
+
+ for file_path in file_list:
+ try:
+ os.remove(file_path)
+ print(f"Deleted file: {file_path}")
+ except OSError as e:
+ print(f"Error deleting file: {file_path}. {str(e)}")
+
+def CRNs2lmdb( dtype,
+ feature_size,
+ feature_name,
+ mean,
+ std,
+ lmdb_dir,
+ lmdb_name
+ ):
+
+ #os.makedirs(os.path.join(lmdb_dir, exist_ok=True))
+ os.makedirs(lmdb_dir, exist_ok=True)
+
+ db_paths = os.path.join(lmdb_dir, "_tmp_data.%04d.lmdb")
+
+ meta_keys = {
+ "dtype" : dtype,
+ "feature_size": feature_size,
+ "feature_name": feature_name,
+ "mean" : mean,
+ "std" : std
+ }
+
+ # Merge LMDB files
+ merge_lmdbs(db_paths, lmdb_dir, lmdb_name)
+ cleanup_lmdb_files(lmdb_dir, "_tmp_data*")
+
+
+
+def write_crns_to_lmdb(mp_args):
+ #pid is idx of workers.
+ db_path, samples, pid, meta_keys = mp_args
+
+ db = lmdb.open(
+ db_path,
+ map_size=1099511627776 * 2,
+ subdir=False,
+ meminit=False,
+ map_async=True,
+ )
+
+ pbar = tqdm(
+ total=len(samples),
+ position=pid,
+ desc=f"Worker {pid}: Writing CRNs Objects into LMDBs",
+ )
+
+ #write indexed samples
+ idx = 0
+ for sample in samples:
+ txn=db.begin(write=True)
+ txn.put(
+ f"{idx}".encode("ascii"),
+ pickle.dumps(sample, protocol=-1),
+ )
+ idx += 1
+ pbar.update(1)
+ txn.commit()
+
+ #write properties
+ txn=db.begin(write=True)
+ txn.put("length".encode("ascii"), pickle.dumps(len(samples), protocol=-1))
+ txn.commit()
+
+ for key, value in meta_keys.items():
+ txn=db.begin(write=True)
+ txn.put(key.encode("ascii"), pickle.dumps(value, protocol=-1))
+ txn.commit()
+
+ db.sync()
+ db.close()
+
+
+def merge_lmdbs(db_paths, out_path, output_file):
+ """
+ merge lmdb files and reordering indexes.
+ """
+ env_out = lmdb.open(
+ os.path.join(out_path, output_file),
+ map_size=1099511627776 * 2,
+ subdir=False,
+ meminit=False,
+ map_async=True,
+ )
+
+
+ idx = 0
+ for db_path in db_paths:
+ env_in = lmdb.open(
+ str(db_path),
+ subdir=False,
+ readonly=True,
+ lock=False,
+ readahead=True,
+ meminit=False,
+ )
+
+ #should set indexes so that properties do not writtent down as well.
+ with env_out.begin(write=True) as txn_out, env_in.begin(write=False) as txn_in:
+ cursor = txn_in.cursor()
+ for key, value in cursor:
+ #write indexed samples
+ try:
+ int(key.decode("ascii"))
+ txn_out.put(
+ f"{idx}".encode("ascii"),
+ value,
+ )
+ idx+=1
+ #print(idx)
+ #write properties
+ except ValueError:
+ txn_out.put(
+ key,
+ value
+ )
+ env_in.close()
+
+ #update length
+ txn_out=env_out.begin(write=True)
+ txn_out.put("length".encode("ascii"), pickle.dumps(idx, protocol=-1))
+ txn_out.commit()
+
+ env_out.sync()
+ env_out.close()
+
+
+
+def write_to_lmdb(new_samples, current_length, lmdb_update, db_path):
+
+ # #pid is idx of workers.
+ # db_path, samples, pid, meta_keys = mp_args
+ db = lmdb.open(
+ db_path,
+ map_size=1099511627776 * 2,
+ subdir=False,
+ meminit=False,
+ map_async=True,
+ )
+
+ pbar = tqdm(
+ total=len(new_samples),
+ desc=f"Adding new samples into LMDBs",
+ )
+
+ #write indexed samples
+ idx = current_length
+ for sample in new_samples:
+ txn=db.begin(write=True)
+ txn.put(
+ f"{idx}".encode("ascii"),
+ pickle.dumps(sample, protocol=-1),
+ )
+ idx += 1
+ pbar.update(1)
+ txn.commit()
+
+ #write properties
+
+ total_length = current_length + len(new_samples)
+
+ txn=db.begin(write=True)
+ txn.put("length".encode("ascii"), pickle.dumps(total_length, protocol=-1))
+ txn.commit()
+
+ #write mean, std, feature_size, feature_name. dtype etc.
+ for key, value in lmdb_update.items():
+ txn=db.begin(write=True)
+ txn.put(key.encode("ascii"), pickle.dumps(value, protocol=-1))
+ txn.commit()
+
+ db.sync()
+ db.close()