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