fennol.training.io

  1import os, io, sys
  2import numpy as np
  3from scipy.spatial.transform import Rotation
  4from collections import defaultdict
  5import pickle
  6import glob
  7from flax import traverse_util
  8from typing import Dict, List, Tuple, Union, Optional, Callable, Sequence
  9from .databases import DBDataset, H5Dataset
 10from ..models.preprocessing import AtomPadding
 11
 12import json
 13import yaml
 14
 15try:
 16    import tomlkit
 17except ImportError:
 18    tomlkit = None
 19
 20try:
 21    from torch.utils.data import DataLoader
 22except ImportError:
 23    raise ImportError(
 24        "PyTorch is required for training. Install the CPU version from https://pytorch.org/get-started/locally/"
 25    )
 26
 27from ..models import FENNIX
 28
 29
 30def load_configuration(config_file: str) -> Dict[str, any]:
 31    if config_file.endswith(".json"):
 32        parameters = json.load(open(config_file))
 33    elif config_file.endswith(".yaml") or config_file.endswith(".yml"):
 34        parameters = yaml.load(open(config_file), Loader=yaml.FullLoader)
 35    elif tomlkit is not None and config_file.endswith(".toml"):
 36        parameters = tomlkit.loads(open(config_file).read())
 37    else:
 38        supported_formats = [".json", ".yaml", ".yml"]
 39        if tomlkit is not None:
 40            supported_formats.append(".toml")
 41        raise ValueError(
 42            f"Unknown config file format. Supported formats: {supported_formats}"
 43        )
 44    return parameters
 45
 46
 47def load_dataset(
 48    dspath: str,
 49    batch_size: int,
 50    rename_refs: dict = {},
 51    infinite_iterator: bool = False,
 52    atom_padding: bool = False,
 53    ref_keys: Optional[Sequence[str]] = None,
 54    split_data_inputs: bool = False,
 55    np_rng: Optional[np.random.Generator] = None,
 56    train_val_split: bool = True,
 57    training_parameters: dict = {},
 58    add_flags: Sequence[str] = ["training"],
 59    fprec: str = "float32",
 60):
 61    """
 62    Load a dataset from a pickle file and return two iterators for training and validation batches.
 63
 64    Args:
 65        training_parameters (dict): A dictionary with the following keys:
 66            - 'dspath': str. Path to the pickle file containing the dataset.
 67            - 'batch_size': int. Number of samples per batch.
 68        rename_refs (list, optional): A list of strings with the names of the reference properties to rename.
 69            Default is an empty list.
 70
 71    Returns:
 72        tuple: A tuple of two infinite iterators, one for training batches and one for validation batches.
 73            For each element in the batch, we expect a "species" key with the atomic numbers of the atoms in the system. Arrays are concatenated along the first axis and the following keys are added to distinguish between the systems:
 74            - 'natoms': np.ndarray. Array with the number of atoms in each system.
 75            - 'batch_index': np.ndarray. Array with the index of the system to which each atom
 76            if the keys "forces", "total_energy", "atomic_energies" or any of the elements in rename_refs are present, the keys are renamed by prepending "true_" to the key name.
 77    """
 78
 79    assert isinstance(
 80        training_parameters, dict
 81    ), "training_parameters must be a dictionary."
 82    assert isinstance(
 83        rename_refs, dict
 84    ), "rename_refs must be a dictionary with the keys to rename."
 85
 86    pbc_training = training_parameters.get("pbc_training", False)
 87    minimum_image = training_parameters.get("minimum_image", False)
 88
 89    coordinates_ref_key = training_parameters.get("coordinates_ref_key", None)
 90
 91    input_keys = [
 92        "species",
 93        "coordinates",
 94        "natoms",
 95        "batch_index",
 96        "total_charge",
 97        "flags",
 98    ]
 99    if pbc_training:
100        input_keys += ["cells", "reciprocal_cells"]
101    if atom_padding:
102        input_keys += ["true_atoms", "true_sys"]
103    if coordinates_ref_key is not None:
104        input_keys += ["system_index", "system_sign"]
105
106    flags = {f: None for f in add_flags}
107    if minimum_image and pbc_training:
108        flags["minimum_image"] = None
109
110    additional_input_keys = set(training_parameters.get("additional_input_keys", []))
111    additional_input_keys_ = set()
112    for key in additional_input_keys:
113        if key not in input_keys:
114            additional_input_keys_.add(key)
115    additional_input_keys = additional_input_keys_
116
117    all_inputs = set(input_keys + list(additional_input_keys))
118
119    extract_all_keys = ref_keys is None
120    if ref_keys is not None:
121        ref_keys = set(ref_keys)
122        ref_keys_ = set()
123        for key in ref_keys:
124            if key not in all_inputs:
125                ref_keys_.add(key)
126
127    random_rotation = training_parameters.get("random_rotation", False)
128    if random_rotation:
129        assert np_rng is not None, "np_rng must be provided for adding noise."
130
131        apply_rotation = {
132            1: lambda x, r: x @ r,
133            -1: lambda x, r: np.einsum("...kn,kj->...jn", x, r),
134            2: lambda x, r: np.einsum("li,...lk,kj->...ij", r, x, r),
135        }
136        valid_rotations = tuple(apply_rotation.keys())
137        rotated_keys = {
138            "coordinates": 1,
139            "forces": 1,
140            "virial_tensor": 2,
141            "stress_tensor": 2,
142            "virial": 2,
143            "stress": 2,
144        }
145        if pbc_training:
146            rotated_keys["cells"] = 1
147        user_rotated_keys = dict(training_parameters.get("rotated_keys", {}))
148        for k, v in user_rotated_keys.items():
149            assert (
150                v in valid_rotations
151            ), f"Invalid rotation type for key {k}. Valid values are {valid_rotations}"
152            rotated_keys[k] = v
153
154        # rotated_vector_keys = set(
155        #     ["coordinates", "forces"]
156        #     + list(training_parameters.get("rotated_vector_keys", []))
157        # )
158        # if pbc_training:
159        #     rotated_vector_keys.add("cells")
160
161        # rotated_tensor_keys = set(
162        #     ["virial_tensor", "stress_tensor", "virial", "stress"]
163        #     + list(training_parameters.get("rotated_tensor_keys", []))
164        # )
165        # assert rotated_vector_keys.isdisjoint(
166        #     rotated_tensor_keys
167        # ), "Rotated vector keys and rotated tensor keys must be disjoint."
168        # rotated_keys = rotated_vector_keys.union(rotated_tensor_keys)
169
170        print(
171            "Applying random rotations to the following keys if present:",
172            list(rotated_keys.keys()),
173        )
174
175        def apply_random_rotations(output, nbatch):
176            euler_angles = np_rng.uniform(0.0, 2 * np.pi, (nbatch, 3))
177            r = [
178                Rotation.from_euler("xyz", euler_angles[i]).as_matrix().T
179                for i in range(nbatch)
180            ]
181            for k, l in rotated_keys.items():
182                if k in output:
183                    for i in range(nbatch):
184                        output[k][i] = apply_rotation[l](output[k][i], r[i])
185
186    else:
187
188        def apply_random_rotations(output, nbatch):
189            pass
190
191    if pbc_training:
192        print("Periodic boundary conditions are active.")
193        length_nopbc = training_parameters.get("length_nopbc", 1000.0)
194
195        def add_cell(d, output):
196            if "cell" not in d:
197                cell = np.asarray(
198                    [
199                        [length_nopbc, 0.0, 0.0],
200                        [0.0, length_nopbc, 0.0],
201                        [0.0, 0.0, length_nopbc],
202                    ],
203                    dtype=fprec,
204                )
205            else:
206                cell = np.asarray(d["cell"], dtype=fprec)
207            output["cells"].append(cell.reshape(1, 3, 3))
208
209    else:
210
211        def add_cell(d, output):
212            if "cell" in d:
213                print(
214                    "Warning: 'cell' found in dataset but not training with pbc. Activate pbc_training to use periodic boundary conditions."
215                )
216
217    if extract_all_keys:
218
219        def add_other_keys(d, output, atom_shift):
220            for k, v in d.items():
221                if k in ("cell", "total_charge"):
222                    continue
223                v_array = np.array(v)
224                # Shift atom number if necessary
225                if k.endswith("_atidx"):
226                    v_array = v_array + atom_shift
227                output[k].append(v_array)
228
229    else:
230
231        def add_other_keys(d, output, atom_shift):
232            output["species"].append(np.asarray(d["species"]))
233            output["coordinates"].append(np.asarray(d["coordinates"], dtype=fprec))
234            for k in additional_input_keys:
235                v_array = np.array(d[k])
236                # Shift atom number if necessary
237                if k.endswith("_atidx"):
238                    v_array = v_array + atom_shift
239                output[k].append(v_array)
240            for k in ref_keys_:
241                v_array = np.array(d[k])
242                # Shift atom number if necessary
243                if k.endswith("_atidx"):
244                    v_array = v_array + atom_shift
245                output[k].append(v_array)
246                if k + "_mask" in d:
247                    output[k + "_mask"].append(np.asarray(d[k + "_mask"]))
248
249    def add_keys(d, output, atom_shift, batch_index):
250        nat = d["species"].shape[0]
251
252        output["natoms"].append(np.asarray([nat]))
253        output["batch_index"].append(np.asarray([batch_index] * nat))
254        if "total_charge" not in d:
255            total_charge = np.asarray(0.0, dtype=fprec)
256        else:
257            total_charge = np.asarray(d["total_charge"], dtype=fprec)
258        output["total_charge"].append(total_charge)
259
260        add_cell(d, output)
261        add_other_keys(d, output, atom_shift)
262
263        return atom_shift + nat
264
265    def collate_fn_(batch):
266        output = defaultdict(list)
267        atom_shift = 0
268        batch_index = 0
269
270        for d in batch:
271            atom_shift = add_keys(d, output, atom_shift, batch_index)
272            batch_index += 1
273
274            if coordinates_ref_key is not None:
275                output["system_index"].append(np.asarray([batch_index - 1]))
276                output["system_sign"].append(np.asarray([1]))
277                if coordinates_ref_key in d:
278                    dref = {**d, "coordinates": d[coordinates_ref_key]}
279                    atom_shift = add_keys(dref, output, atom_shift, batch_index)
280                    output["system_index"].append(np.asarray([batch_index - 1]))
281                    output["system_sign"].append(np.asarray([-1]))
282                    batch_index += 1
283
284        nbatch_ = len(output["natoms"])
285        apply_random_rotations(output,nbatch_)
286
287        # Stack and concatenate the arrays
288        for k, v in output.items():
289            if v[0].ndim == 0:
290                v = np.stack(v)
291            else:
292                v = np.concatenate(v, axis=0)
293            if np.issubdtype(v.dtype, np.floating):
294                v = v.astype(fprec)
295            output[k] = v
296
297        if "cells" in output and pbc_training:
298            output["reciprocal_cells"] = np.linalg.inv(output["cells"])
299
300        # Rename necessary keys
301        # for key in rename_refs:
302        #     if key in output:
303        #         output["true_" + key] = output.pop(key)
304        for kold, knew in rename_refs.items():
305            assert (
306                knew not in output
307            ), f"Cannot rename key {kold} to {knew}. Key {knew} already present."
308            if kold in output:
309                output[knew] = output.pop(kold)
310
311        output["flags"] = flags
312        return output
313
314    collate_layers_train = [collate_fn_]
315    collate_layers_valid = [collate_fn_]
316
317    ### collate preprocessing
318    # add noise to the training data
319    noise_sigma = training_parameters.get("noise_sigma", None)
320    if noise_sigma is not None:
321        assert isinstance(noise_sigma, dict), "noise_sigma should be a dictionary"
322
323        for sigma in noise_sigma.values():
324            assert sigma >= 0, "Noise sigma should be a positive number"
325
326        print("Adding noise to the training data:")
327        for key, sigma in noise_sigma.items():
328            print(f"  - {key} with sigma = {sigma}")
329
330        assert np_rng is not None, "np_rng must be provided for adding noise."
331
332        def collate_with_noise(batch):
333            for key, sigma in noise_sigma.items():
334                if key in batch and sigma > 0:
335                    batch[key] += np_rng.normal(0, sigma, batch[key].shape).astype(
336                        batch[key].dtype
337                    )
338            return batch
339
340        collate_layers_train.append(collate_with_noise)
341
342    if atom_padding:
343        padder = AtomPadding(add_sys=training_parameters.get("padder_add_sys", 0))
344        padder_state = padder.init()
345
346        def collate_with_padding(batch):
347            padder_state_up, output = padder(padder_state, batch)
348            padder_state.update(padder_state_up)
349            return output
350
351        collate_layers_train.append(collate_with_padding)
352        collate_layers_valid.append(collate_with_padding)
353
354    if split_data_inputs:
355
356        # input_keys += additional_input_keys
357        # input_keys = set(input_keys)
358        print("Input keys:", all_inputs)
359        print("Ref keys:", ref_keys)
360
361        def collate_split(batch):
362            inputs = {}
363            refs = {}
364            for k, v in batch.items():
365                if k in all_inputs:
366                    inputs[k] = v
367                if k in ref_keys:
368                    refs[k] = v
369                if k.endswith("_mask") and k[:-5] in ref_keys:
370                    refs[k] = v
371            return inputs, refs
372
373        collate_layers_train.append(collate_split)
374        collate_layers_valid.append(collate_split)
375
376    ### apply all collate preprocessing
377    if len(collate_layers_train) == 1:
378        collate_fn_train = collate_layers_train[0]
379    else:
380
381        def collate_fn_train(batch):
382            for layer in collate_layers_train:
383                batch = layer(batch)
384            return batch
385
386    if len(collate_layers_valid) == 1:
387        collate_fn_valid = collate_layers_valid[0]
388    else:
389
390        def collate_fn_valid(batch):
391            for layer in collate_layers_valid:
392                batch = layer(batch)
393            return batch
394
395    # dspath = training_parameters.get("dspath", None)
396    print(f"Loading dataset from {dspath}...", end="")
397    # print(f"   the following keys will be renamed if present : {rename_refs}")
398    sharded_training = False
399    if dspath.endswith(".db"):
400        dataset = {}
401        if train_val_split:
402            dataset["training"] = DBDataset(dspath, table="training")
403            dataset["validation"] = DBDataset(dspath, table="validation")
404        else:
405            dataset = DBDataset(dspath)
406    elif dspath.endswith(".h5") or dspath.endswith(".hdf5"):
407        dataset = {}
408        if train_val_split:
409            dataset["training"] = H5Dataset(dspath, table="training")
410            dataset["validation"] = H5Dataset(dspath, table="validation")
411        else:
412            dataset = H5Dataset(dspath)
413    elif dspath.endswith(".pkl") or dspath.endswith(".pickle"):
414        with open(dspath, "rb") as f:
415            dataset = pickle.load(f)
416        if not train_val_split and isinstance(dataset, dict):
417            dataset = dataset["training"]
418    elif os.path.isdir(dspath):
419        if train_val_split:
420            dataset = {}
421            with open(dspath + "/validation.pkl", "rb") as f:
422                dataset["validation"] = pickle.load(f)
423        else:
424            dataset = None
425
426        shard_files = sorted(glob.glob(dspath + "/training_*.pkl"))
427        nshards = len(shard_files)
428        if nshards == 0:
429            raise ValueError("No dataset shards found.")
430        elif nshards == 1:
431            with open(shard_files[0], "rb") as f:
432                if train_val_split:
433                    dataset["training"] = pickle.load(f)
434                else:
435                    dataset = pickle.load(f)
436        else:
437            print(f"Found {nshards} dataset shards.")
438            sharded_training = True
439
440    else:
441        raise ValueError(
442            f"Unknown dataset format. Supported formats: '.db', '.h5', '.pkl', '.pickle'"
443        )
444    print(" done.")
445
446    ### BUILD DATALOADERS
447    # batch_size = training_parameters.get("batch_size", 16)
448    shuffle = training_parameters.get("shuffle_dataset", True)
449    if train_val_split:
450        dataloader_validation = DataLoader(
451            dataset["validation"],
452            batch_size=batch_size,
453            shuffle=shuffle,
454            collate_fn=collate_fn_valid,
455        )
456
457    if sharded_training:
458
459        def iterate_sharded_dataset():
460            indices = np.arange(nshards)
461            if shuffle:
462                assert np_rng is not None, "np_rng must be provided for shuffling."
463                np_rng.shuffle(indices)
464            for i in indices:
465                filename = shard_files[i]
466                print(f"# Loading dataset shard from {filename}...", end="")
467                with open(filename, "rb") as f:
468                    dataset = pickle.load(f)
469                print(" done.")
470                dataloader = DataLoader(
471                    dataset,
472                    batch_size=batch_size,
473                    shuffle=shuffle,
474                    collate_fn=collate_fn_train,
475                )
476                for batch in dataloader:
477                    yield batch
478
479        class DataLoaderSharded:
480            def __iter__(self):
481                return iterate_sharded_dataset()
482
483        dataloader_training = DataLoaderSharded()
484    else:
485        dataloader_training = DataLoader(
486            dataset["training"] if train_val_split else dataset,
487            batch_size=batch_size,
488            shuffle=shuffle,
489            collate_fn=collate_fn_train,
490        )
491
492    if not infinite_iterator:
493        if train_val_split:
494            return dataloader_training, dataloader_validation
495        return dataloader_training
496
497    def next_batch_factory(dataloader):
498        while True:
499            for batch in dataloader:
500                yield batch
501
502    training_iterator = next_batch_factory(dataloader_training)
503    if train_val_split:
504        validation_iterator = next_batch_factory(dataloader_validation)
505        return training_iterator, validation_iterator
506    return training_iterator
507
508
509def load_model(
510    parameters: Dict[str, any],
511    model_file: Optional[str] = None,
512    rng_key: Optional[str] = None,
513) -> FENNIX:
514    """
515    Load a FENNIX model from a file or create a new one.
516
517    Args:
518        parameters (dict): A dictionary of parameters for the model.
519        model_file (str, optional): The path to a saved model file to load.
520
521    Returns:
522        FENNIX: A FENNIX model object.
523    """
524    print_model = parameters["training"].get("print_model", False)
525    if model_file is None:
526        model_file = parameters.get("model_file", None)
527    if model_file is not None and os.path.exists(model_file):
528        model = FENNIX.load(model_file, use_atom_padding=False)
529        if print_model:
530            print(model.summarize())
531        print(f"Restored model from '{model_file}'.")
532    else:
533        assert (
534            rng_key is not None
535        ), "rng_key must be specified if model_file is not provided."
536        model_params = parameters["model"]
537        if isinstance(model_params, str):
538            assert os.path.exists(
539                model_params
540            ), f"Model file '{model_params}' not found."
541            model = FENNIX.load(model_params, use_atom_padding=False)
542            print(f"Restored model from '{model_params}'.")
543        else:
544            model = FENNIX(**model_params, rng_key=rng_key, use_atom_padding=False)
545        if print_model:
546            print(model.summarize())
547    return model
548
549
550def copy_parameters(variables, variables_ref, params):
551    def merge_params(full_path_, v, v_ref):
552        full_path = "/".join(full_path_[1:]).lower()
553        status = (False, "")
554        for path in params:
555            if full_path.startswith(path.lower()) and len(path) > len(status[1]):
556                status = (True, path)
557        return v_ref if status[0] else v
558
559    flat = traverse_util.flatten_dict(variables, keep_empty_nodes=False)
560    flat_ref = traverse_util.flatten_dict(variables_ref, keep_empty_nodes=False)
561    return traverse_util.unflatten_dict(
562        {
563            k: merge_params(k, v, flat_ref[k]) if k in flat_ref else v
564            for k, v in flat.items()
565        }
566    )
567
568
569class TeeLogger(object):
570    def __init__(self, name):
571        self.file = io.TextIOWrapper(open(name, "wb"), write_through=True)
572        self.stdout = None
573
574    def __del__(self):
575        self.close()
576
577    def write(self, data):
578        self.file.write(data)
579        self.stdout.write(data)
580        self.flush()
581
582    def close(self):
583        self.file.close()
584
585    def flush(self):
586        self.file.flush()
587
588    def bind_stdout(self):
589        if isinstance(sys.stdout, TeeLogger):
590            raise ValueError("stdout already bound to a Tee instance.")
591        if self.stdout is not None:
592            raise ValueError("stdout already bound.")
593        self.stdout = sys.stdout
594        sys.stdout = self
595
596    def unbind_stdout(self):
597        if self.stdout is None:
598            raise ValueError("stdout is not bound.")
599        sys.stdout = self.stdout