qupled 1.3.2__cp312-cp312-macosx_14_0_arm64.whl → 1.3.4__cp312-cp312-macosx_14_0_arm64.whl

qupled/database.py

@@ -10,8 +10,6 @@ import sqlalchemy as sql
 from sqlalchemy.dialects.sqlite import insert as sqlite_insert
 import blosc2
 
-from . import mpi
-
 
 class DataBaseHandler:
     """
@@ -42,10 +40,9 @@ class DataBaseHandler:
         SUCCESS = "SUCCESS"
         FAILED = "FAILED"
 
-    INT_TO_RUN_STATUS = {
-        0: RunStatus.SUCCESS,
-        1: RunStatus.FAILED,
-    }
+    class ConflictMode(Enum):
+        FAIL = "FAIL"
+        UPDATE = "UPDATE"
 
     def __init__(self, database_name: str | None = None):
         """
@@ -76,7 +73,6 @@ class DataBaseHandler:
         self.result_table = self._build_results_table()
         self.run_id: int | None = None
 
-    @mpi.MPI.run_only_on_root
     def insert_run(self, inputs):
         """
         Inserts a new run into the database by storing the provided inputs and results.
@@ -91,7 +87,6 @@ class DataBaseHandler:
         self._insert_run(inputs, self.RunStatus.RUNNING)
         self.insert_inputs(inputs.__dict__)
 
-    @mpi.MPI.run_only_on_root
     def insert_inputs(self, inputs: dict[str, any]):
         """
         Inserts input data into the database for the current run.
@@ -114,8 +109,11 @@ class DataBaseHandler:
             sql_mapping = lambda value: (self._to_json(value))
             self._insert_from_dict(self.input_table, inputs, sql_mapping)
 
-    @mpi.MPI.run_only_on_root
-    def insert_results(self, results: dict[str, any]):
+    def insert_results(
+        self,
+        results: dict[str, any],
+        conflict_mode: ConflictMode = ConflictMode.FAIL,
+    ):
         """
         Inserts the given results into the database table associated with this instance.
 
@@ -130,7 +128,9 @@ class DataBaseHandler:
         """
         if self.run_id is not None:
             sql_mapping = lambda value: (self._to_bytes(value))
-            self._insert_from_dict(self.result_table, results, sql_mapping)
+            self._insert_from_dict(
+                self.result_table, results, sql_mapping, conflict_mode
+            )
 
     def inspect_runs(self) -> list[dict[str, any]]:
         """
@@ -149,26 +149,27 @@ class DataBaseHandler:
         rows = self._execute(statement).mappings().all()
         return [{key: row[key] for key in row.keys()} for row in rows]
 
-    def update_run_status(self, status: int) -> None:
+    def update_run_status(self, status: RunStatus) -> None:
         """
-        Updates the status of a run in the database.
+        Update the status of a run in the database.
 
         Args:
-            status (int): The new status code to update the run with. If the status
-                          code is not found in the INT_TO_RUN_STATUS mapping, the
-                          status will default to RunStatus.FAILED.
+            status (RunStatus): The new status to set for the run.
 
         Returns:
             None
+
+        Notes:
+            This method updates the status of the run identified by `self.run_id` in the run table.
+            If `self.run_id` is None, no update is performed.
         """
         if self.run_id is not None:
-            new_status = self.INT_TO_RUN_STATUS.get(status, self.RunStatus.FAILED)
             statement = (
                 sql.update(self.run_table)
                 .where(
                     self.run_table.c[self.TableKeys.PRIMARY_KEY.value] == self.run_id
                 )
-                .values({self.TableKeys.STATUS.value: new_status.value})
+                .values({self.TableKeys.STATUS.value: status.value})
             )
             self._execute(statement)
 
@@ -241,8 +242,6 @@ class DataBaseHandler:
         sql_mapping = lambda value: (self._from_bytes(value))
         return self._get(self.result_table, run_id, names, sql_mapping)
 
-    @mpi.MPI.synchronize_ranks
-    @mpi.MPI.run_only_on_root
     def delete_run(self, run_id: int) -> None:
         """
         Deletes a run entry from the database based on the provided run ID.
@@ -390,16 +389,15 @@ class DataBaseHandler:
             sql.PrimaryKeyConstraint(
                 self.TableKeys.RUN_ID.value, self.TableKeys.NAME.value
             ),
+            sql.Index(f"idx_{table_name}_run_id", self.TableKeys.RUN_ID.value),
+            sql.Index(f"idx_{table_name}_name", self.TableKeys.NAME.value),
         )
         self._create_table(table)
         return table
 
-    @mpi.MPI.synchronize_ranks
-    @mpi.MPI.run_only_on_root
     def _create_table(self, table):
         table.create(self.engine, checkfirst=True)
 
-    @mpi.MPI.run_only_on_root
     def _insert_run(self, inputs: any, status: RunStatus):
         """
         Inserts a new run entry into the database.
@@ -454,11 +452,15 @@ class DataBaseHandler:
         def _set_pragma(dbapi_connection, connection_record):
             cursor = dbapi_connection.cursor()
             cursor.execute("PRAGMA foreign_keys=ON")
+            cursor.execute("PRAGMA journal_mode=WAL")
             cursor.close()
 
-    @mpi.MPI.run_only_on_root
     def _insert_from_dict(
-        self, table, data: dict[str, any], sql_mapping: Callable[[any], any]
+        self,
+        table,
+        data: dict[str, any],
+        sql_mapping: Callable[[any], any],
+        conflict_mode: ConflictMode = ConflictMode.FAIL,
     ) -> None:
         """
         Inserts data into a specified table by mapping values through a provided SQL mapping function.
@@ -473,10 +475,46 @@ class DataBaseHandler:
         """
         for name, value in data.items():
             if mapped_value := sql_mapping(value):
-                self._insert(table, name, mapped_value)
+                self._insert(table, name, mapped_value, conflict_mode)
+
+    def _insert(
+        self,
+        table: sql.Table,
+        name: str,
+        value: any,
+        conflict_mode: ConflictMode = ConflictMode.FAIL,
+    ):
+        """
+        Inserts a record into the specified SQL table with the given name and value, handling conflicts according to the specified mode.
+        Args:
+            table (sql.Table): The SQLAlchemy table object where the record will be inserted.
+            name (str): The name/key associated with the value to insert.
+            value (any): The value to be inserted into the table.
+            conflict_mode (ConflictMode, optional): Specifies how to handle conflicts on unique constraints.
+                Defaults to ConflictMode.FAIL. If set to ConflictMode.UPDATE, existing records with the same
+                run_id and name will be updated with the new value.
+        Returns:
+            None
+        Raises:
+            Any exceptions raised by the underlying database execution.
+        """
+        data = {
+            self.TableKeys.RUN_ID.value: self.run_id,
+            self.TableKeys.NAME.value: name,
+            self.TableKeys.VALUE.value: value,
+        }
+        statement = sqlite_insert(table).values(data)
+        if conflict_mode == self.ConflictMode.UPDATE:
+            statement = statement.on_conflict_do_update(
+                index_elements=[
+                    self.TableKeys.RUN_ID.value,
+                    self.TableKeys.NAME.value,
+                ],
+                set_={self.TableKeys.VALUE.value: value},
+            )
+        self._execute(statement)
 
-    @mpi.MPI.run_only_on_root
-    def _insert(self, table: sql.Table, name: str, value: any):
+    def _insert_with_update(self, table: sql.Table, name: str, value: any):
         """
         Inserts a record into the specified SQL table or updates it if a conflict occurs.
 

qupled/esa.py

@@ -2,13 +2,17 @@ from __future__ import annotations
 
 from . import hf
 from . import native
+from . import serialize
 
 
-class ESA(hf.HF):
+class Solver(hf.Solver):
     """
     Class used to solve the ESA scheme.
     """
 
+    # Native classes used to solve the scheme
+    native_scheme_cls = native.ESA
+
     def __init__(self):
         super().__init__()
         self.results: hf.Result = hf.Result()
@@ -16,12 +20,14 @@ class ESA(hf.HF):
         self.native_scheme_cls = native.ESA
 
 
+@serialize.serializable_dataclass
 class Input(hf.Input):
     """
     Class used to manage the input for the :obj:`qupled.esa.ESA` class.
     """
 
-    def __init__(self, coupling: float, degeneracy: float):
-        super().__init__(coupling, degeneracy)
-        # Undocumented default values
-        self.theory = "ESA"
+    theory: str = "ESA"
+
+
+if __name__ == "__main__":
+    Solver.run_mpi_worker(Input, hf.Result)

qupled/hf.py

@@ -1,17 +1,34 @@
 from __future__ import annotations
 
+import json
+
+from dataclasses import field
+from pathlib import Path
+
 import numpy as np
 
 from . import database
 from . import mpi
 from . import native
+from . import serialize
+from . import timer
 
 
-class HF:
+class Solver:
     """
     Class used to solve the HF scheme.
     """
 
+    # Mapping of native scheme status to run status in the database
+    NATIVE_TO_RUN_STATUS = {
+        0: database.DataBaseHandler.RunStatus.SUCCESS,
+        1: database.DataBaseHandler.RunStatus.FAILED,
+    }
+
+    # Native classes used to solve the scheme
+    native_scheme_cls = native.HF
+    native_inputs_cls = native.Input
+
     def __init__(self):
         self.inputs: Input = None
         """The inputs used to solve the scheme. Default = ``None``"""
@@ -19,8 +36,6 @@ class HF:
         """The results obtained by solving the scheme"""
         # Undocumented properties
         self.db_handler = database.DataBaseHandler()
-        self.native_scheme_cls = native.HF
-        self.native_inputs_cls = native.Input
         self.native_scheme_status = None
 
     @property
@@ -33,8 +48,7 @@ class HF:
         """
         return self.db_handler.run_id
 
-    @mpi.MPI.record_time
-    @mpi.MPI.synchronize_ranks
+    @timer.timer
     def compute(self, inputs: Input):
         """
         Solves the scheme and saves the results.
@@ -47,7 +61,6 @@ class HF:
         self._compute_native()
         self._save()
 
-    @mpi.MPI.run_only_on_root
     def compute_rdf(self, rdf_grid: np.ndarray = None):
         """
         Computes the radial distribution function (RDF) using the provided RDF grid.
@@ -61,7 +74,8 @@ class HF:
         if self.results is not None:
             self.results.compute_rdf(rdf_grid)
             self.db_handler.insert_results(
-                {"rdf": self.results.rdf, "rdf_grid": self.results.rdf_grid}
+                {"rdf": self.results.rdf, "rdf_grid": self.results.rdf_grid},
+                conflict_mode=database.DataBaseHandler.ConflictMode.UPDATE,
             )
 
     def _add_run_to_database(self):
@@ -76,6 +90,19 @@ class HF:
         self.inputs.database_info.run_id = self.run_id
 
     def _compute_native(self):
+        """
+        Determines whether to execute the native computation in parallel or serial mode.
+
+        Checks if MPI (Message Passing Interface) is available and if the number of requested processes
+        is greater than one. If both conditions are met, runs the computation in parallel; otherwise,
+        runs it in serial mode.
+        """
+        if native.uses_mpi:
+            self._compute_native_mpi()
+        else:
+            self._compute_native_serial()
+
+    def _compute_native_serial(self):
         """
         Computes the native representation of the inputs and processes the results.
 
@@ -91,7 +118,32 @@ class HF:
         self.native_scheme_status = scheme.compute()
         self.results.from_native(scheme)
 
-    @mpi.MPI.run_only_on_root
+    def _compute_native_mpi(self):
+        """
+        Executes a native MPI computation workflow.
+
+        This method performs the following steps:
+        1. Writes the necessary input files for the MPI computation using `mpi.write_inputs`.
+        2. Launches the MPI execution by calling `mpi.launch_mpi_execution` with the current module and the specified number of processes.
+        3. Reads the computation results using `mpi.read_results` and assigns them to `self.results`.
+        4. Cleans up any temporary files generated during the computation with `mpi.clean_files`.
+        """
+        mpi.write_inputs(self.inputs)
+        mpi.launch_mpi_execution(self.__module__, self.inputs.processes)
+        self.native_scheme_status = mpi.read_status()
+        self.results = mpi.read_results(type(self.results))
+        mpi.clean_files()
+
+    @classmethod
+    def run_mpi_worker(cls, InputCls, ResultCls):
+        inputs = mpi.read_inputs(InputCls)
+        native_inputs = cls.native_inputs_cls()
+        inputs.to_native(native_inputs)
+        scheme = cls.native_scheme_cls(native_inputs)
+        status = scheme.compute()
+        mpi.write_results(scheme, ResultCls)
+        mpi.write_status(scheme, status)
+
     def _save(self):
         """
         Saves the current state and results to the database.
@@ -99,56 +151,53 @@ class HF:
         This method updates the run status in the database using the current
         native scheme status and inserts the results into the database.
         """
-        self.db_handler.update_run_status(self.native_scheme_status)
+        run_status = self.NATIVE_TO_RUN_STATUS.get(
+            self.native_scheme_status, database.DataBaseHandler.RunStatus.FAILED
+        )
+        self.db_handler.update_run_status(run_status)
         self.db_handler.insert_results(self.results.__dict__)
 
 
+@serialize.serializable_dataclass
 class Input:
     """
     Class used to store the inputs for the :obj:`qupled.hf.HF` class.
     """
 
-    def __init__(self, coupling: float, degeneracy: float):
-        """
-        Initialize the base class with the given parameters.
-
-        Parameters:
-            coupling (float): Coupling parameter.
-            degeneracy (float): Degeneracy parameter.
-        """
-        self.chemical_potential: list[float] = [-10.0, 10.0]
-        """Initial guess for the chemical potential. Default = ``[-10, 10]``"""
-        self.coupling: float = coupling
-        """Coupling parameter."""
-        self.cutoff: float = 10.0
-        """Cutoff for the wave-vector grid. Default =  ``10.0``"""
-        self.degeneracy: float = degeneracy
-        """Degeneracy parameter."""
-        self.frequency_cutoff: float = 10.0
-        """Cutoff for the frequency (applies only in the ground state). Default =  ``10.0``"""
-        self.integral_error: float = 1.0e-5
-        """Accuracy (relative error) in the computation of integrals. Default = ``1.0e-5``"""
-        self.integral_strategy: str = "full"
-        """
-        Scheme used to solve two-dimensional integrals
-        allowed options include:
+    coupling: float
+    """Coupling parameter."""
+    degeneracy: float
+    """Degeneracy parameter."""
+    chemical_potential: list[float] = field(default_factory=lambda: [-10.0, 10.0])
+    """Initial guess for the chemical potential. Default = ``[-10, 10]``"""
+    cutoff: float = 10.0
+    """Cutoff for the wave-vector grid. Default =  ``10.0``"""
+    frequency_cutoff: float = 10.0
+    """Cutoff for the frequency (applies only in the ground state). Default =  ``10.0``"""
+    integral_error: float = 1.0e-5
+    """Accuracy (relative error) in the computation of integrals. Default = ``1.0e-5``"""
+    integral_strategy: str = "full"
+    """
+    Scheme used to solve two-dimensional integrals
+    allowed options include:
 
-        - full: the inner integral is evaluated at arbitrary points selected automatically by the quadrature rule
+    - full: the inner integral is evaluated at arbitrary points selected automatically by the quadrature rule
 
-        - segregated: the inner integral is evaluated on a fixed grid that depends on the integrand that is being processed
+    - segregated: the inner integral is evaluated on a fixed grid that depends on the integrand that is being processed
 
-        Segregated is usually faster than full but it could become
-        less accurate if the fixed points are not chosen correctly. Default =  ``'full'``
-        """
-        self.matsubara: int = 128
-        """Number of Matsubara frequencies. Default = ``128``"""
-        self.resolution: float = 0.1
-        """Resolution of the wave-vector grid. Default =  ``0.1``"""
-        self.threads: int = 1
-        """Number of OMP threads for parallel calculations. Default =  ``1``"""
-        # Undocumented default values
-        self.theory: str = "HF"
-        self.database_info: DatabaseInfo = DatabaseInfo()
+    Segregated is usually faster than full but it could become
+    less accurate if the fixed points are not chosen correctly. Default =  ``'full'``
+    """
+    matsubara: int = 128
+    """Number of Matsubara frequencies. Default = ``128``"""
+    resolution: float = 0.1
+    """Resolution of the wave-vector grid. Default =  ``0.1``"""
+    threads: int = 1
+    """Number of OMP threads for parallel calculations. Default =  ``1``"""
+    processes: int = 1
+    """Number of MPI processes for parallel calculations. Default =  ``1``"""
+    theory: str = "HF"
+    database_info: DatabaseInfo = field(default_factory=lambda: DatabaseInfo())
 
     def to_native(self, native_input: any):
         """
@@ -176,28 +225,28 @@ class Input:
                 setattr(native_input, attr, value_to_set)
 
 
+@serialize.serializable_dataclass
 class Result:
     """
     Class used to store the results for the :obj:`qupled.hf.HF` class.
     """
 
-    def __init__(self):
-        self.idr: np.ndarray = None
-        """Ideal density response"""
-        self.lfc: np.ndarray = None
-        """Local field correction"""
-        self.rdf: np.ndarray = None
-        """Radial distribution function"""
-        self.rdf_grid: np.ndarray = None
-        """Radial distribution function grid"""
-        self.sdr: np.ndarray = None
-        """Static density response"""
-        self.ssf: np.ndarray = None
-        """Static structure factor"""
-        self.uint: float = None
-        """Internal energy"""
-        self.wvg: np.ndarray = None
-        """Wave-vector grid"""
+    idr: np.ndarray = None
+    """Ideal density response"""
+    lfc: np.ndarray = None
+    """Local field correction"""
+    rdf: np.ndarray = None
+    """Radial distribution function"""
+    rdf_grid: np.ndarray = None
+    """Radial distribution function grid"""
+    sdr: np.ndarray = None
+    """Static density response"""
+    ssf: np.ndarray = None
+    """Static structure factor"""
+    uint: float = None
+    """Internal energy"""
+    wvg: np.ndarray = None
+    """Wave-vector grid"""
 
     def from_native(self, native_scheme: any):
         """
@@ -210,10 +259,11 @@ class Result:
             - Only attributes that exist in both the current object and the native_scheme object will be updated.
             - Attributes with a value of `None` in the native_scheme object will not overwrite the current object's attributes.
         """
-        for attr in self.__dict__.keys():
+        for attr in self.__dataclass_fields__:
             if hasattr(native_scheme, attr):
                 value = getattr(native_scheme, attr)
-                setattr(self, attr, value) if value is not None else None
+                valid_value = value is not None and not callable(value)
+                setattr(self, attr, value) if valid_value else None
 
     def compute_rdf(self, rdf_grid: np.ndarray | None = None):
         """
@@ -234,18 +284,18 @@ class Result:
             self.rdf = native.compute_rdf(self.rdf_grid, self.wvg, self.ssf)
 
 
+@serialize.serializable_dataclass
 class DatabaseInfo:
     """
     Class used to store the database information passed to the native code.
     """
 
-    def __init__(self):
-        self.name: str = database.DataBaseHandler.DEFAULT_DATABASE_NAME
-        """Database name"""
-        self.run_id: int = None
-        """ID of the run in the database"""
-        self.run_table_name: str = database.DataBaseHandler.RUN_TABLE_NAME
-        """Name of the table used to store the runs in the database"""
+    name: str = database.DataBaseHandler.DEFAULT_DATABASE_NAME
+    """Database name"""
+    run_id: int = None
+    """ID of the run in the database"""
+    run_table_name: str = database.DataBaseHandler.RUN_TABLE_NAME
+    """Name of the table used to store the runs in the database"""
 
     def to_native(self) -> native.DatabaseInfo:
         """
@@ -261,3 +311,14 @@ class DatabaseInfo:
             if value is not None:
                 setattr(native_database_info, attr, value)
         return native_database_info
+
+    @classmethod
+    def from_dict(cls, d):
+        obj = cls.__new__(cls)
+        for key, value in d.items():
+            setattr(obj, key, value)
+        return obj
+
+
+if __name__ == "__main__":
+    Solver.run_mpi_worker(Input, Result)