qclab 0.3.0a7__py3-none-any.whl

qclab/__init__.py

@@ -0,0 +1,13 @@
+"""
+This module contains the QC Lab package initialization.
+"""
+
+from qclab.simulation import Simulation
+from qclab.variable import Variable
+from qclab.data import Data
+from qclab.model import Model
+from qclab.algorithm import Algorithm
+from qclab.constants import Constants
+from qclab.utils import configure_memory_logger
+
+configure_memory_logger()

qclab/algorithm.py

@@ -0,0 +1,43 @@
+"""
+This module contains the Algorithm class.
+"""
+
+import copy
+from qclab.constants import Constants
+
+
+class Algorithm:
+    """
+    Algorithm class for defining and executing algorithm recipes.
+    """
+
+    def __init__(self, default_settings=None, settings=None):
+        if settings is None:
+            settings = {}
+        if default_settings is None:
+            default_settings = {}
+        # Merge default settings with user-provided settings.
+        settings = {**default_settings, **settings}
+        # Construct a Constants object to hold settings.
+        self.settings = Constants()
+        # Put settings from the dictionary into the Constants object.
+        for key, val in settings.items():
+            setattr(self.settings, key, val)
+        # Copy the recipes and output variables to ensure they are not shared
+        # across instances.
+        self.initialization_recipe = copy.deepcopy(self.initialization_recipe)
+        self.update_recipe = copy.deepcopy(self.update_recipe)
+        self.collect_recipe = copy.deepcopy(self.collect_recipe)
+
+    initialization_recipe = []
+    update_recipe = []
+    collect_recipe = []
+
+    def execute_recipe(self, sim, state, parameters, recipe):
+        """
+        Carry out the given recipe for the simulation by running
+        each task in the recipe.
+        """
+        for func in recipe:
+            state, parameters = func(sim, state, parameters)
+        return state, parameters

qclab/algorithms/__init__.py

@@ -0,0 +1,8 @@
+"""
+This module imports the algorithm classes to qclab.algorithms.
+"""
+
+from qclab.algorithms.mean_field import MeanField
+from qclab.algorithms.fewest_switches_surface_hopping import (
+    FewestSwitchesSurfaceHopping,
+)

qclab/algorithms/fewest_switches_surface_hopping.py

@@ -0,0 +1,124 @@
+"""
+This module contains the FewestSwitchesSurfaceHopping algorithm class.
+"""
+
+from functools import partial
+from qclab.algorithm import Algorithm
+from qclab import tasks
+
+
+class FewestSwitchesSurfaceHopping(Algorithm):
+    """
+    Fewest switches surface hopping algorithm class.
+    """
+
+    def __init__(self, settings=None):
+        if settings is None:
+            settings = {}
+        self.default_settings = {
+            "fssh_deterministic": False,
+            "gauge_fixing": "sign_overlap",
+            "use_gauge_field_force": False,
+        }
+        super().__init__(self.default_settings, settings)
+
+    initialization_recipe = [
+        tasks.initialize_variable_objects,
+        tasks.initialize_norm_factor,
+        tasks.initialize_branch_seeds,
+        tasks.initialize_z,
+        tasks.update_h_quantum,
+        partial(
+            tasks.diagonalize_matrix,
+            matrix_name="h_quantum",
+            eigvals_name="eigvals",
+            eigvecs_name="eigvecs",
+        ),
+        partial(
+            tasks.gauge_fix_eigs,
+            gauge_fixing="phase_der_couple",
+            eigvecs_previous_name="eigvecs",
+        ),
+        partial(tasks.copy_in_state, copy_name="eigvecs_previous", orig_name="eigvecs"),
+        partial(
+            tasks.basis_transform_vec,
+            input_vec_name="wf_db",
+            basis_name="eigvecs",
+            output_vec_name="wf_adb",
+            adb_to_db=False,
+        ),
+        tasks.initialize_random_values_fssh,
+        tasks.initialize_active_surface,
+        tasks.initialize_dm_adb_0_fssh,
+        tasks.update_act_surf_wf,
+    ]
+
+    update_recipe = [
+        partial(tasks.copy_in_state, copy_name="eigvecs_previous", orig_name="eigvecs"),
+        # Begin RK4 integration steps.
+        tasks.update_classical_forces,
+        partial(
+            tasks.update_quantum_classical_forces,
+            wf_db_name="act_surf_wf",
+            wf_changed=True,
+        ),
+        tasks.update_z_rk4_k1,
+        partial(tasks.update_classical_forces, z_name="z_1"),
+        partial(
+            tasks.update_quantum_classical_forces,
+            wf_db_name="act_surf_wf",
+            z_name="z_1",
+            wf_changed=False,
+        ),
+        tasks.update_z_rk4_k2,
+        partial(tasks.update_classical_forces, z_name="z_2"),
+        partial(
+            tasks.update_quantum_classical_forces,
+            wf_db_name="act_surf_wf",
+            z_name="z_2",
+            wf_changed=False,
+        ),
+        tasks.update_z_rk4_k3,
+        partial(tasks.update_classical_forces, z_name="z_3"),
+        partial(
+            tasks.update_quantum_classical_forces,
+            wf_db_name="act_surf_wf",
+            z_name="z_3",
+            wf_changed=False,
+        ),
+        tasks.update_z_rk4_k4,
+        # End RK4 integration steps.
+        tasks.update_wf_db_eigs,
+        tasks.update_h_quantum,
+        partial(
+            tasks.diagonalize_matrix,
+            matrix_name="h_quantum",
+            eigvals_name="eigvals",
+            eigvecs_name="eigvecs",
+        ),
+        tasks.gauge_fix_eigs,
+        partial(
+            tasks.basis_transform_vec,
+            input_vec_name="wf_db",
+            basis_name="eigvecs",
+            output_vec_name="wf_adb",
+            adb_to_db=False,
+        ),
+        tasks.update_hop_probs_fssh,
+        tasks.update_hop_inds_fssh,
+        tasks.update_hop_vals_fssh,
+        tasks.update_z_hop_fssh,
+        tasks.update_act_surf_hop_fssh,
+        tasks.update_act_surf_wf,
+    ]
+
+    collect_recipe = [
+        tasks.update_t,
+        tasks.update_dm_db_fssh,
+        tasks.update_quantum_energy_fssh,
+        tasks.update_classical_energy_fssh,
+        tasks.collect_t,
+        tasks.collect_dm_db,
+        tasks.collect_quantum_energy,
+        tasks.collect_classical_energy,
+    ]

qclab/algorithms/mean_field.py

@@ -0,0 +1,68 @@
+"""
+This module contains the MeanField algorithm class.
+"""
+
+from functools import partial
+from qclab.algorithm import Algorithm
+from qclab import tasks
+
+
+class MeanField(Algorithm):
+    """
+    Mean-field dynamics algorithm class.
+    """
+
+    def __init__(self, settings=None):
+        if settings is None:
+            settings = {}
+        self.default_settings = {}
+        super().__init__(self.default_settings, settings)
+
+    initialization_recipe = [
+        tasks.initialize_variable_objects,
+        tasks.initialize_norm_factor,
+        tasks.initialize_z,
+        tasks.update_h_quantum,
+    ]
+
+    update_recipe = [
+        # Begin RK4 integration steps.
+        partial(tasks.update_classical_forces, z_name="z"),
+        tasks.update_quantum_classical_forces,
+        tasks.update_z_rk4_k1,
+        partial(tasks.update_classical_forces, z_name="z_1"),
+        partial(
+            tasks.update_quantum_classical_forces,
+            z_name="z_1",
+            wf_changed=False,
+        ),
+        tasks.update_z_rk4_k2,
+        partial(tasks.update_classical_forces, z_name="z_2"),
+        partial(
+            tasks.update_quantum_classical_forces,
+            z_name="z_2",
+            wf_changed=False,
+        ),
+        tasks.update_z_rk4_k3,
+        partial(tasks.update_classical_forces, z_name="z_3"),
+        partial(
+            tasks.update_quantum_classical_forces,
+            z_name="z_3",
+            wf_changed=False,
+        ),
+        tasks.update_z_rk4_k4,
+        # End RK4 integration steps.
+        tasks.update_wf_db_rk4,
+        tasks.update_h_quantum,
+    ]
+
+    collect_recipe = [
+        tasks.update_t,
+        tasks.update_dm_db_mf,
+        tasks.update_quantum_energy,
+        tasks.update_classical_energy,
+        tasks.collect_t,
+        tasks.collect_dm_db,
+        tasks.collect_classical_energy,
+        tasks.collect_quantum_energy,
+    ]

qclab/constants.py

@@ -0,0 +1,37 @@
+"""
+This module contains the Constants class.
+"""
+
+
+class Constants:
+    """
+    Constants class to define constants and trigger updates when they change.
+    """
+
+    def __init__(self, update_function=None):
+        self._updating = False
+        self._init_complete = False
+        self._update_function = update_function
+
+    def __setattr__(self, name, value):
+        """
+        Overrides attribute setting to call the update function after the attribute is
+        changed, preventing recursion.
+        """
+        super().__setattr__(name, value)
+        if not self._updating and name not in {
+            "_updating",
+            "_update_function",
+            "_init_complete",
+        }:
+            if self._init_complete:
+                self._updating = True
+                if self._update_function is not None:
+                    self._update_function()
+                self._updating = False
+
+    def get(self, name, default=None):
+        """
+        Get the value of an attribute, returning default if it does not exist.
+        """
+        return getattr(self, name, default)

qclab/data.py

@@ -0,0 +1,186 @@
+"""
+This module contains the Data class.
+"""
+
+import logging
+import numpy as np
+from qclab.utils import DISABLE_H5PY
+
+if not DISABLE_H5PY:
+    import h5py
+
+logger = logging.getLogger(__name__)
+
+
+class Data:
+    """
+    Data class for handling the collection of data during a simulation.
+    """
+
+    def __init__(self, seeds=None):
+        if seeds is None:
+            seeds = np.array([], dtype=int)
+        self.data_dict = {"seed": seeds, "norm_factor": 0}
+        # Store log messages captured during a simulation run. This attribute is
+        # populated by the drivers when they return the Data object.
+        self.log = ""
+
+    def add_output_to_data_dict(self, sim, state, t_ind):
+        """
+        Add data to the output dictionary ``self.data_dict``.
+
+        .. rubric:: Args
+        sim: Simulation
+            The simulation object containing settings and parameters.
+        state: Variable
+            The state object containing the current simulation state.
+        t_ind: int
+            The current time index in the simulation.
+        """
+        # Check if the norm_factor is zero. If it is, save it from the state object.
+        if self.data_dict["norm_factor"] == 0:
+            if not hasattr(state, "norm_factor"):
+                logger.critical(
+                    "The state object does not have a norm_factor attribute. "
+                    "This is required to normalize the data."
+                )
+                raise ValueError(
+                    "The state object does not have a norm_factor attribute."
+                )
+            logger.info(
+                "Setting norm_factor to %s from state object.", state.norm_factor
+            )
+            self.data_dict["norm_factor"] = state.norm_factor
+        for key, val in state.output_dict.items():
+            if not key in self.data_dict:
+                # If the key is not in the data_dict, initialize it with zeros.
+                self.data_dict[key] = np.zeros(
+                    (len(sim.settings.t_collect), *np.shape(val)[1:]), dtype=val.dtype
+                )
+                logger.info(
+                    "Initializing data_dict[%s] with shape %s.",
+                    key,
+                    self.data_dict[key].shape,
+                )
+            # Store the data in the data_dict at the correct time index.
+            self.data_dict[key][t_ind // sim.settings.dt_collect_n] = (
+                np.sum(val, axis=0) / self.data_dict["norm_factor"]
+            )
+
+    def add_data(self, new_data):
+        """
+        Add data from ``new_data`` to the output dictionary ``self.data_dict``.
+
+        .. rubric:: Args
+        new_data: Data
+            A Data instance containing the new data to merge.
+        """
+        new_norm_factor = (
+            new_data.data_dict["norm_factor"] + self.data_dict["norm_factor"]
+        )
+        for key, val in new_data.data_dict.items():
+            if key == "seed":
+                self.data_dict[key] = np.concatenate(
+                    (self.data_dict[key], val.flatten()), axis=0
+                )
+            elif key != "norm_factor":
+                if key in self.data_dict:
+                    self.data_dict[key] = (
+                        self.data_dict[key] * self.data_dict["norm_factor"]
+                        + val * new_data.data_dict["norm_factor"]
+                    ) / new_norm_factor
+                else:
+                    self.data_dict[key] = val
+        self.data_dict["norm_factor"] = new_norm_factor
+        # Append any log messages stored in new_data to this instance's log.
+        if getattr(new_data, "log", ""):
+            self.log += new_data.log
+
+    def save(self, filename, disable_h5py=DISABLE_H5PY):
+        """
+        Save the data to disk with file name ``filename``.
+
+        If h5py is available the data is stored as an HDF5 archive; otherwise
+        each variable is saved using numpy.savez.
+
+        .. rubric:: Args
+        filename : str
+            The file name to save the data to.
+        disable_h5py : bool, default: qclab.utils.DISABLE_H5PY
+            If True, h5py is not used even if available.
+        """
+        if disable_h5py:
+            np.savez(filename, log=self.log, **self.data_dict)
+        else:
+            with h5py.File(filename, "w") as h5file:
+                self._recursive_save(h5file, "/", self.data_dict)
+                h5file.attrs["log"] = self.log
+
+    def load(self, filename, disable_h5py=DISABLE_H5PY):
+        """
+        Load a Data object from ``filename``.
+
+        .. rubric:: Args
+        filename : str
+            The file name to load the data from.
+        disable_h5py : bool, default: qclab.utils.DISABLE_H5PY
+            If True, h5py is not used even if available.
+
+        .. rubric:: Returns
+        Data : Data
+            The loaded Data object.
+        """
+        new_data = Data()
+        if disable_h5py:
+            loaded = np.load(filename, allow_pickle=True)
+            new_data.data_dict = {key: loaded[key] for key in loaded.files if key != "log"}
+            new_data.log = str(loaded.get("log", ""))
+            self.add_data(new_data)
+            return self
+        with h5py.File(filename, "r") as h5file:
+            new_data._recursive_load(h5file, "/", new_data.data_dict)
+            new_data.log = h5file.attrs["log"]
+            self.add_data(new_data)
+        return self
+
+    def _recursive_save(self, h5file, path, dict):
+        """
+        Recursively save dictionary contents to an HDF5 group.
+
+        .. rubric:: Args
+        h5file : h5py.File
+            The HDF5 file object to save data into.
+        path : str
+            The current path in the HDF5 file.
+        dict : dict
+            The dictionary to save.
+        """
+        for key, item in dict.items():
+            if isinstance(item, type(dict)):
+                self._recursive_save(h5file, path + key + "/", item)
+            else:
+                try:
+                    h5file.create_dataset(path + key, data=np.asarray(item))
+                except TypeError as exc:
+                    raise ValueError(
+                        f"Cannot save {key} with unsupported type {type(item)}."
+                    ) from exc
+
+    def _recursive_load(self, h5file, path, dict):
+        """
+        Recursively load dictionary contents from an HDF5 group.
+
+        .. rubric:: Args
+        h5file : h5py.File
+            The HDF5 file object to load data from.
+        path : str
+            The current path in the HDF5 file.
+        dict : dict
+            The dictionary to populate with loaded data.
+        """
+        for key, item in h5file[path].items():
+            if isinstance(item, h5py._hl.dataset.Dataset):
+                dict[key] = item[()]
+            elif isinstance(item, h5py._hl.group.Group):
+                dict[key] = {}
+                self._recursive_load(h5file, path + key + "/", dict[key])

qclab/dynamics/__init__.py

@@ -0,0 +1,10 @@
+"""
+This module imports the dynamics drivers to qclab.dynamics.
+"""
+
+from qclab.dynamics.parallel_driver_multiprocessing import (
+    parallel_driver_multiprocessing,
+)
+from qclab.dynamics.parallel_driver_mpi import parallel_driver_mpi
+from qclab.dynamics.serial_driver import serial_driver
+from qclab.dynamics.dynamics import run_dynamics

qclab/dynamics/dynamics.py

@@ -0,0 +1,51 @@
+"""
+This module contains the dynamics core.
+"""
+
+import numpy as np
+from tqdm import tqdm
+
+
+def run_dynamics(sim, state, parameters, data):
+    """
+    Dynamics core for QC Lab.
+
+    .. rubric:: Args
+    sim: Simulation
+        The simulation object containing the model, algorithm, and settings.
+    state: Variable
+        The state variable object containing the simulation seeds.
+    parameters: Variable
+        The parameters variable object containing any additional parameters.
+    data: Data
+        The data object for collecting output data.
+
+    .. rubric:: Returns
+    data: Data
+        The updated data object containing collected output data.
+    """
+    # Define an update iterator using tqdm if progress_bar is True.
+    t_update_iterator = sim.settings.t_update_n
+    if getattr(sim.settings, "progress_bar", True):
+        t_update_iterator = tqdm(t_update_iterator)
+
+    # Iterate over each time step.
+    for sim.t_ind in t_update_iterator:
+        if sim.t_ind == 0:
+            # Execute initialization recipe.
+            state, parameters = sim.algorithm.execute_recipe(
+                sim, state, parameters, sim.algorithm.initialization_recipe
+            )
+        # Detect collect timesteps.
+        if np.mod(sim.t_ind, sim.settings.dt_collect_n) == 0:
+            # Calculate output variables.
+            state, parameters = sim.algorithm.execute_recipe(
+                sim, state, parameters, sim.algorithm.collect_recipe
+            )
+            # Collect totals in output dictionary.
+            data.add_output_to_data_dict(sim, state, sim.t_ind)
+        # Execute update recipe.
+        state, parameters = sim.algorithm.execute_recipe(
+            sim, state, parameters, sim.algorithm.update_recipe
+        )
+    return data

qclab/dynamics/parallel_driver_mpi.py

@@ -0,0 +1,129 @@
+"""
+This module contains the parallel MPI driver.
+"""
+
+import logging
+import copy
+import numpy as np
+import qclab.dynamics as dynamics
+from qclab.utils import get_log_output
+from qclab import Variable, Data
+
+logger = logging.getLogger(__name__)
+
+
+def parallel_driver_mpi(sim, seeds=None, data=None, num_tasks=None):
+    """
+    Parallel driver for the dynamics core using the mpi4py library.
+
+    .. rubric:: Args
+    sim: Simulation
+        The simulation object containing the model, algorithm, initial state, and settings.
+    seeds: ndarray, optional
+        An array of integer seeds for the trajectories. If None, seeds will be
+        generated automatically.
+    data: Data, optional
+        A Data object for collecting output data. If None, a new Data object
+        will be created.
+    num_tasks: int, optional
+        The number of tasks to use for parallel processing. If None, the
+        number of available tasks will be used.
+
+    .. rubric:: Returns
+    data: Data
+        The updated Data object containing collected output data.
+    """
+    # First initialize the model constants.
+    sim.model.initialize_constants()
+    try:
+        from mpi4py import MPI
+    except ImportError:
+        raise ImportError(
+            "The package mpi4py is required for the parallel_driver_mpi driver."
+        ) from None
+    except Exception as e:
+        raise RuntimeError(f"An error occurred when importing mpi4py: {e}") from None
+    if seeds is None:
+        offset = 0
+        if data is not None:
+            if len(data.data_dict["seed"]) > 0:
+                offset = np.max(data.data_dict["seed"]) + 1
+        else:
+            data = Data()
+        seeds = offset + np.arange(sim.settings.num_trajs, dtype=int)
+        num_trajs = sim.settings.num_trajs
+    else:
+        num_trajs = len(seeds)
+        logger.warning(
+            "Setting sim.settings.num_trajs to the number of provided seeds."
+        )
+        sim.settings.num_trajs = num_trajs
+        if data is None:
+            data = Data()
+    comm = MPI.COMM_WORLD
+    rank = comm.Get_rank()
+    if num_tasks is None:
+        size = comm.Get_size()
+    else:
+        size = num_tasks
+    # Determine the number of simulations required to execute the total number
+    # of trajectories.
+    if num_trajs % sim.settings.batch_size == 0:
+        num_sims = num_trajs // sim.settings.batch_size
+    else:
+        num_sims = num_trajs // sim.settings.batch_size + 1
+    batch_seeds_list = (
+        np.zeros((num_sims * sim.settings.batch_size), dtype=int) + np.nan
+    )
+    batch_seeds_list[:num_trajs] = seeds
+    batch_seeds_list = batch_seeds_list.reshape((num_sims, sim.settings.batch_size))
+    # Split the simulations into chunks for each MPI process.
+    chunk_inds = np.linspace(0, num_sims, size + 1, dtype=int)
+    start = chunk_inds[rank]
+    end = chunk_inds[rank + 1]
+    chunk_size = end - start
+    # Create the input data for each local simulation.
+    sim.initialize_timesteps()
+    local_input_data = [
+        (
+            copy.deepcopy(sim),
+            Variable(
+                {
+                    "seed": batch_seeds_list[n][~np.isnan(batch_seeds_list[n])].astype(
+                        int
+                    )
+                }
+            ),
+            Variable(),
+            Data(batch_seeds_list[n][~np.isnan(batch_seeds_list[n])].astype(int)),
+        )
+        for n in np.arange(num_sims)[start:end]
+    ]
+    # Set the batch size for each local simulation.
+    for i in range(chunk_size):
+        local_input_data[i][0].settings.batch_size = len(local_input_data[i][1].seed)
+    # Execute the local simulations.
+    local_results = [dynamics.run_dynamics(*x) for x in local_input_data]
+    comm.Barrier()
+    # Collect results sequentially on rank 0.
+    tag_data, tag_done = 1, 2
+    if rank == 0:
+        for result in local_results:
+            data.add_data(result)
+        remaining = size - 1
+        status = MPI.Status()
+        while remaining:
+            msg = comm.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
+            if status.Get_tag() == tag_done:
+                remaining -= 1
+            else:
+                data.add_data(msg)
+    else:
+        for result in local_results:
+            comm.send(result, dest=0, tag=tag_data)
+        comm.send(None, dest=0, tag=tag_done)
+    # Collect logs from all ranks and attach combined output on root rank.
+    gathered_logs = comm.gather(get_log_output(), root=0)
+    if rank == 0:
+        data.log = "".join(log for log in gathered_logs if log)
+    return data