### Install OpenFerro from Source

Source: https://openferro.readthedocs.io/en/latest/_sources/installation.rst.txt

Clone the OpenFerro repository from GitHub, navigate into the directory, and install OpenFerro using pip. This method installs the library from its source code.

```bash
git clone https://github.com/salinelake/OpenFerro.git
cd OpenFerro
pip install .
```

--------------------------------

### Install JAX (CPU-only)

Source: https://openferro.readthedocs.io/en/latest/installation.html

Install the CPU-only version of JAX, which is required for OpenFerro. Ensure JAX version 0.4.0 or later is installed.

```bash
pip install -U jax

```

--------------------------------

### Install JAX (GPU with CUDA 12)

Source: https://openferro.readthedocs.io/en/latest/_sources/installation.rst.txt

Install JAX with CUDA 12 support for GPU acceleration. Ensure you have a compatible NVIDIA GPU and CUDA toolkit installed.

```bash
pip install -U "jax[cuda12]"
```

--------------------------------

### Install JAX (CPU-only)

Source: https://openferro.readthedocs.io/en/latest/_sources/installation.rst.txt

Install the CPU-only version of JAX. This is suitable for Linux, macOS, and Windows systems without a compatible NVIDIA GPU.

```bash
pip install -U jax
```

--------------------------------

### Initialize Reporters

Source: https://openferro.readthedocs.io/en/latest/api.html

Initializes all added reporters before starting the simulation.

```python
sim.initialize_reporters()
```

--------------------------------

### Verify OpenFerro Installation

Source: https://openferro.readthedocs.io/en/latest/_sources/installation.rst.txt

Import the OpenFerro library in a Python interpreter to confirm a successful installation. This is a basic check to ensure the library is accessible.

```python
import openferro as of
```

--------------------------------

### Get All Fields

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/system.html

Returns a list of all fields currently present in the system.

```python
all_fields = system.get_all_fields()

```

--------------------------------

### _openferro.parallelism.DeviceMesh.__init__

Source: https://openferro.readthedocs.io/en/latest/api.html

Initializes the single-host multi-device parallelism by getting the mesh of the devices. It allows specifying devices, number of rows, and number of columns for the device mesh.

```APIDOC
## _openferro.parallelism.DeviceMesh.__init__

### Description
Initialize the single-host multi-device parallelism. Get the mesh of the devices.

### Parameters
#### Path Parameters
* **devices** (array-like, optional) - List of devices to use. If None, uses all available devices
* **num_rows** (int, optional) - Number of rows in device mesh. If None, automatically determined
* **num_cols** (int, optional) - Number of columns in device mesh. If None, automatically determined

### Raises
**ValueError** - If only one device is available If num_rows * num_cols does not match number of devices
```

--------------------------------

### Initialize Simulation Reporters

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/simulation.html

Initializes all reporters added to the simulation. This method should be called before starting the simulation.

```python
self.initialize_reporters()
```

--------------------------------

### Initialize Integrator Base Class

Source: https://openferro.readthedocs.io/en/latest/api.html

Initializes the base Integrator class with a specified time step. This is a fundamental setup for any integration process.

```python
def __init__(self, dt):
    pass
```

--------------------------------

### Ewald Energy Engine Setup

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/engine/ewald.html

Sets up the Ewald energy engine for calculating dipole-dipole interactions. It pre-computes reciprocal space grid and coefficients for efficiency. This function is intended for use within the OpenFerro engine.

```python
def get_dipole_dipole_ewald_high_memory_usage(latt):
    """
    Returns the function to calculate the energy of dipole-dipole interaction.
    Implemented according to Sec.5.3 of 
    "Wang, D., et al. "Ewald summation for ferroelectric perovksites with charges and dipoles." Computational Materials Science 162 (2019): 314-321."
    """
    l1, l2, l3 = latt.size
    a1, a2, a3 = latt.latt_vec
    a1 = a1[0]
    a2 = a2[1]
    a3 = a3[2]
    ref_volume = a1 * a2 * a3 * l1 * l2 * l3
    a = jnp.array([a1 , a2 , a3 ])
    b = 2 * jnp.pi / a
    bmax = jnp.max(b)
    amin = 2 * np.pi / bmax
    alpha = 5 / amin
    gcut = 2 * np.pi * alpha
    sigma = 1.0 / alpha / jnp.sqrt(2.0)   ## the ewald sigma parameter
    
    ## get coefficients
    coef_ksum = 1 / 2.0 / ref_volume / Constants.epsilon0
    coef_rsum = 1 / 2.0 / jnp.pi / Constants.epsilon0 * alpha**3 / 3.0 / jnp.sqrt(jnp.pi) 

    ## get reriprocal space grid
    n1 = int(gcut / b[0])
    n2 = int(gcut / b[1])
    n3 = int(gcut / b[2])
    ng1, ng2, ng3 = l1*n1, l2*n2, l3*n3
    G_grid = jnp.stack( jnp.meshgrid(
        jnp.arange(-ng1, ng1) / l1 * b[0], 
        jnp.arange(-ng2, ng2) / l2 * b[1], 
        jnp.arange(-ng3, ng3) / l3 * b[2], 
        indexing='ij'), axis=-1)   # (2*ng1, 2*ng2, 2*ng3, 3)
    G_grid = jnp.roll(G_grid, shift=(-ng1, -ng2, -ng3), axis=(0, 1, 2))  # move gamma point to (0,0,0)
    G_grid = G_grid.reshape(2*n1, l1, 2*n2, l2, 2*n3, l3, 3)     
    G_grid = G_grid.transpose(1,3,5,0,2,4,6).reshape(l1,l2,l3,-1,3)  # (l1, l2, l3, 8*n1*n2*n3, 3)

    ## get coefficients for reciprocal space sum
    Uk_coef = jnp.exp( - 0.5 * sigma**2 * jnp.sum(G_grid**2, axis=-1) ) / jnp.sum(G_grid**2, axis=-1)   # (l1, l2, l3, 8*n1*n2*n3)
    Uk_coef = Uk_coef.at[0,0,0,0].set(0.0)   # mute Gamma point
    def energy_engine(field, parameters):
        Z = parameters['Z_star']
        epsilon_inf = parameters['epsilon_inf']

        ## calculate reciprocal space sum
        F_fft3 = jnp.fft.fftn(field, axes=(0,1,2))  # (l1, l2, l3, 3)
        Uk_squared  = jnp.sum( F_fft3.real[:,:,:,None,:] * G_grid, axis=-1)**2
        Uk_squared += jnp.sum( F_fft3.imag[:,:,:,None,:] * G_grid, axis=-1)**2   # (l1, l2, l3, 8*n1*n2*n3)
        ewald_ksum = coef_ksum * jnp.sum(Uk_coef * Uk_squared)

        ## calculate real space sum
        ewald_rsum = - coef_rsum * jnp.sum(field**2)
        return (ewald_ksum + ewald_rsum) * Z**2 / epsilon_inf
    return energy_engine
```

--------------------------------

### Initialize Langevin NVT Simulation

Source: https://openferro.readthedocs.io/en/latest/api.html

Sets up a constant volume, temperature (NVT) molecular dynamics simulation using Langevin dynamics. Requires a System object.

```python
nvt_langevin_sim = SimulationNVTLangevin(system)
```

--------------------------------

### Initialize NVE Simulation

Source: https://openferro.readthedocs.io/en/latest/api.html

Sets up a microcanonical (NVE) molecular dynamics simulation. Requires a System object.

```python
nve_sim = SimulationNVE(system)
```

--------------------------------

### LLSIBIntegrator __init__ method

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/integrator/llg.html

Initializes the LLSIBIntegrator with simulation parameters. It sets up the time step, damping constant, gyromagnetic ratio, and convergence criteria for iterative solvers. The update functions are JIT compiled.

```python
def __init__(self, dt, alpha, gamma=None, max_iter=10, tol=1e-6):
        super().__init__(dt)
        ## get the renormalized gyromagnetic ratio for simulating LLG equation in Landau-Lifshitz form
        self.alpha = alpha
        if gamma is None:
            self.gamma = Constants.electron_gyro_ratio / (1 + alpha**2)
        else:
            self.gamma = gamma / (1 + alpha**2)
        self.max_iter = max_iter  
        self.tol = tol
        self.update_x = jit(self._update_x)
        self.update_b = jit(self._update_b)
```

--------------------------------

### System.__init__

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/system.html

Initializes a System object with a given lattice.

```APIDOC
## Method __init__

### Description
Initializes a System object.

### Parameters
- **lattice** (BravaisLattice3D): The lattice of the system.
```

--------------------------------

### Initialize Energy Minimization

Source: https://openferro.readthedocs.io/en/latest/api.html

Initializes an energy minimization simulation. Requires a System object and optional convergence parameters.

```python
md_min = MDMinimize(system, max_iter=500, tol=1e-6)
```

--------------------------------

### Get Field Variance

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/field.html

Calculates the variance of the field over the lattice dimensions.

```python
@property
    def var(self):
        """
        Calculate the variance of the field over the lattice.

        Returns
        -------
        array_like
            Variance of field
        """
        return jnp.var(self.get_values(), axis=[i for i in range(self.ldim)])
```

--------------------------------

### Get Field Mean

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/field.html

Calculates the mean value of the field over the lattice dimensions.

```python
@property
    def mean(self):
        """
        Calculate the average of the field over the lattice.

        Returns
        -------
        array_like
            Mean value of field
        """
        return jnp.mean(self.get_values(), axis=[i for i in range(self.ldim)])
```

--------------------------------

### Get Interaction Parameters

Source: https://openferro.readthedocs.io/en/latest/api.html

Retrieves the current parameters of the interaction as a jax numpy array.

```python
openferro.interaction.interaction_base.get_parameters()
```

--------------------------------

### Create OpenFerro System

Source: https://openferro.readthedocs.io/en/latest/_sources/logic.rst.txt

Initializes an OpenFerro system with a defined lattice. The system object will hold fields and global variables.

```python
system = of.System(lat)
```

--------------------------------

### Get All Non-SO3 Fields

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/system.html

Filters and returns all fields that are not instances of the SO3 field type.

```python
non_so3_fields = system.get_all_non_SO3_fields()

```

--------------------------------

### RingPolymerSystem Initialization

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/system.html

Initializes a RingPolymerSystem, inheriting from System. Currently raises NotImplementedError as it is not yet implemented.

```python
class RingPolymerSystem(System):
    """
    A class to define a ring polymer system for path-integral molecular dynamics simulations. 
    """

    def __init__(self, lattice, nbeads=1):
        super().__init__(lattice)
        self.nbeads = nbeads
        raise NotImplementedError("Ring polymer system is not implemented yet.")
```

--------------------------------

### Get All SO3 Fields

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/system.html

Filters and returns only the fields that are instances of the SO3 field type.

```python
so3_fields = system.get_all_SO3_fields()

```

--------------------------------

### Run Simulation Steps

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/simulation.html

Initializes reporters and runs the simulation for a specified number of steps, performing a single step in each iteration. Includes a sanity check to ensure integrators are set for all fields.

```python
## sanity check
        for field in self.all_fields:
            if field.integrator is None:
                raise ValueError('Please set the integrator for the field %s before running the simulation' % type(field))
        ## run the simulation
        self.initialize_reporters()
        for i in range(nsteps):
            self._step(profile=profile)
            self.step_reporters()
```

--------------------------------

### Get Field Velocity

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/field.html

Retrieves the velocity of the field. Raises a ValueError if velocity has not been set.

```python
def get_velocity(self):
    if self._velocity is None:
        raise ValueError("Velocity is not set")
    else:
        return self._velocity
```

--------------------------------

### Initialize NPT Langevin Simulation

Source: https://openferro.readthedocs.io/en/latest/api.html

Initializes a constant pressure, constant temperature (NPT) molecular dynamics simulation using Langevin dynamics. Requires a System object and an optional pressure value.

```python
openferro.simulation.SimulationNPTLangevin(_system_ , _pressure =0.0_)
```

--------------------------------

### Get Field Mass

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/field.html

Retrieves the mass of the field. Raises a ValueError if mass has not been set.

```python
def get_mass(self):
    if self._mass is None:
        raise ValueError("Mass is not set")
    else:
        return self._mass
```

--------------------------------

### Initialize NPT Langevin Simulation

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/simulation.html

Initializes an NPT molecular dynamics simulation using Langevin dynamics, setting the external pressure.

```python
    def __init__(self, system, pressure=0.0):
        super().__init__(system)
        ## set pressure
        self.pressure = pressure
        pV_param = self.system.get_interaction_by_ID('pV').get_parameters()
        pV_param_new = [pressure * Constants.bar, pV_param[1]]
        self.system.get_interaction_by_ID('pV').set_parameters(pV_param_new)
        ## get all fields, including the global strain field
        self.SO3_fields = self.system.get_all_SO3_fields()
        self.non_SO3_fields = self.system.get_all_non_SO3_fields()
        self.all_fields = self.SO3_fields + self.non_SO3_fields
        self.nfields = len(self.all_fields)
```

--------------------------------

### Get Field Force

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/field.html

Retrieves the force acting on the field. Raises a ValueError if force has not been set.

```python
def get_force(self):
    if self._force is None:
        raise ValueError("Force do not exist")
    else:
        return self._force
```

--------------------------------

### MDMinimize Initialization

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/simulation.html

Initializes the MDMinimize simulation for energy minimization. Sets maximum iterations and force tolerance.

```python
MDMinimize(system, max_iter=100, tol=1e-5 )
```

--------------------------------

### Get Field Values

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/field.html

Retrieves the current values of the field. Raises a ValueError if values have not been set.

```python
def get_values(self):
    """
    Get the values of the field.

    Returns
    -------
    array_like
        Values of the field

    Raises
    ------
    ValueError
        If field has no values set
    """
    if self._values is None:
        raise ValueError("Field has no values. Set values before getting them.")
    else:
        return self._values
```

--------------------------------

### Initialize System with Lattice

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/system.html

Initializes a System object with a given lattice. This sets up the system's periodicity and internal dictionaries for fields and interactions.

```python
system = System(lattice)

```

--------------------------------

### Get Total Number of Sites

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/lattice.html

Calculates the total number of sites in the lattice by multiplying the dimensions.

```python
@property
    def nsites(self):
        return jnp.prod(self.size)
```

--------------------------------

### Initialize LLSIBLangevinIntegrator

Source: https://openferro.readthedocs.io/en/latest/api.html

Initializes the integrator with time step, temperature, damping coefficient, and optional gyromagnetic ratio and iteration parameters. Used for setting up stochastic spin dynamics simulations.

```python
LLSIBLangevinIntegrator(_dt_ , _temp_ , _alpha_ , _gamma =None_, _max_iter =10_, _tol =1e-06_)
```

--------------------------------

### Get All Interactions (Python)

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/system.html

Retrieves a dictionary containing all self, mutual, and triple interactions currently defined in the system.

```python
@property
def interaction_dict(self):
        """
        Get all interactions in the system.

        Returns
        -------
        dict
            All interactions in the system
        """
        return {**self._self_interaction_dict, **self._mutual_interaction_dict, **self._triple_interaction_dict}
```

--------------------------------

### Get Magnitude of SO(3) Field

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/field.html

Retrieves the current magnitude of the SO(3) field. Raises an error if the magnitude has not been set.

```python
magnitude = field_so3.get_magnitude()

```

--------------------------------

### LLSIBLangevinIntegrator Constructor

Source: https://openferro.readthedocs.io/en/latest/api.html

Initializes the LLSIBLangevinIntegrator with specified parameters for simulating spin dynamics.

```APIDOC
## LLSIBLangevinIntegrator Constructor

### Description
Initializes the LLSIBLangevinIntegrator with specified parameters for simulating spin dynamics.

### Parameters

- **_dt_**: float - Time step size.
- **_temp_**: float - Temperature.
- **_alpha_**: float - Damping parameter.
- **_gamma_**: float, optional - Gyromagnetic ratio. Defaults to None.
- **_max_iter_**: int, optional - Maximum number of iterations for the solver. Defaults to 10.
- **_tol_**: float, optional - Tolerance for the solver. Defaults to 1e-06.
```

--------------------------------

### LangevinIntegrator Initialization

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/integrator/md.html

Initializes the LangevinIntegrator with time step, temperature, and relaxation time. It pre-computes constants for efficiency.

```python
def __init__(self, dt, temp, tau):
        super().__init__(dt)
        self.temp = temp
        self.kbT = Constants.kb * temp
        self.tau = tau
        self.gamma = 1.0 / tau
        self.z1 = jnp.exp(-dt * self.gamma)
        self.z2 = jnp.sqrt(1 - jnp.exp(-2 * dt * self.gamma))
        self.step_p = jit(self._step_p)
        self.step_x = jit(self._step_x)
        self.step_t = jit(self._step_t)
```

--------------------------------

### Get Field by ID

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/system.html

Retrieves a specific field from the system using its unique ID. Raises a ValueError if the ID does not exist.

```python
field = system.get_field_by_ID('my_field_id')

```

--------------------------------

### Get Lattice Vectors

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/lattice.html

Stacks the primitive vectors a1, a2, and a3 into a 3x3 array representing the lattice vectors.

```python
@property
    def latt_vec(self):
        return jnp.stack([self.a1, self.a2, self.a3], axis=0)
```

--------------------------------

### Initialize NVE Simulation in OpenFerro

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/simulation.html

Initializes the NVE (microcanonical) simulation by setting up SO3 and non-SO3 fields, excluding the global strain field.

```python
super().__init__(system)
## get all fields, excluding the global strain field
self.SO3_fields = self.system.get_all_SO3_fields()
self.non_SO3_fields = [field for field in self.system.get_all_non_SO3_fields() if not isinstance(field, GlobalStrain)]
self.all_fields = self.SO3_fields + self.non_SO3_fields
self.nfields = len(self.all_fields)
```

--------------------------------

### Interaction Base Classes

Source: https://openferro.readthedocs.io/en/latest/api.html

Base classes for defining interactions between fields, including setting and getting parameters, and calculating energy and forces.

```APIDOC
## Class openferro.interaction.interaction_base

### Description
The base class to specify the interaction between fields. Each interaction is associated with a term in the Hamiltonian.

### Parameters
* **parameters** (None, optional) - Parameters for the interaction

### Methods
#### __init__(self, parameters=None)
Initializes the interaction_base object.

#### set_parameters(self, parameters)
Sets the parameters of the interaction.

Parameters:
* **parameters** (array_like) - The parameters of the interaction

Raises:
* **ValueError** - If parameters is not a numpy array, list, or jax array

#### get_parameters(self)
Gets the parameters of the interaction.

Returns:
The parameters of the interaction (jax.numpy.ndarray)

#### set_energy_engine(self, energy_engine, enable_jit=True)
Sets the energy engine of the interaction.

Parameters:
* **energy_engine** (callable) - The energy engine of the interaction. It should take the values of the fields as input and return the energy as output.
* **enable_jit** (bool, optional) - Whether to enable jit for the energy engine, by default True

#### calc_energy(self)
Calculates the energy of the interaction.

#### calc_force(self)
Calculates the force of the interaction.
```

```APIDOC
## Class openferro.interaction.self_interaction

### Description
A class to specify the self-interaction of a field.

### Parameters
* **field_ID** (str) - Identifier for the field
* **parameters** (array_like, optional) - Parameters for the interaction, by default None

### Methods
#### __init__(self, field_ID, parameters=None)
Initializes the self_interaction object.

#### create_force_engine(self, enable_jit=True)
Derives the force engine of the interaction from the energy engine through automatic differentiation.

Parameters:
* **enable_jit** (bool, optional) - Whether to enable jit for the force engine, by default True

Raises:
* **ValueError** - If energy engine is not set

#### calc_energy(self, field)
Calculates the energy of the interaction for a given field.

Parameters:
* **field** (Field) - The field to calculate the energy

Returns:
The energy of the interaction (float)

#### calc_force(self, field)
Calculates the force of the interaction for a given field.

Parameters:
* **field** (Field) - The field to calculate the force

Returns:
The gradient of the energy with respect to the field. It has the same shape as the field (jax.numpy.ndarray)
```

```APIDOC
## Class openferro.interaction.mutual_interaction

### Description
A class to specify the mutual interaction between two fields.

### Parameters
* **field_1_ID** (str) - Identifier for the first field
* **field_2_ID** (str) - Identifier for the second field
* **parameters** (array_like, optional) - Parameters for the interaction, by default None

### Methods
#### __init__(self, field_1_ID, field_2_ID, parameters=None)
Initializes the mutual_interaction object.

#### create_force_engine(self, enable_jit=True)
Derives the force engine of the interaction from the energy engine through automatic differentiation.

Parameters:
* **enable_jit** (bool, optional) - Whether to enable jit for the force engine, by default True

Raises:
* **ValueError** - If energy engine is not set

#### calc_energy(self, field1, field2)
Calculates the energy of the interaction for a given pair of fields.

Parameters:
* **field1** (Field) - The first field
* **field2** (Field) - The second field
```

--------------------------------

### LLSIBLangevinIntegrator __init__ method

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/integrator/llg.html

Initializes the LLSIBLangevinIntegrator for stochastic LLG simulations. It inherits from LLSIBIntegrator and adds temperature-dependent parameters for thermal noise calculation.

```python
def __init__(self, dt, temp, alpha, gamma=None, max_iter=10, tol=1e-6):
        super().__init__(dt, alpha, gamma, max_iter, tol)
        """
        Args:
            dt: the time step
            temp: the temperature
            alpha: the Gilbert damping constant
            gamma: the gyromagnetic ratio
        """
        self.kbT = Constants.kb * temp
        self.D_base = self.alpha/(1+self.alpha**2) * self.kbT / self.gamma  
        self.noise_max = jnp.sqrt(2* jnp.abs(jnp.log(self.dt)))
```

--------------------------------

### Get Interaction by ID (Python)

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/system.html

Retrieves a specific interaction from the system using its unique ID. Raises a ValueError if the interaction does not exist.

```python
def get_interaction_by_ID(self, interaction_ID):
        """
        Get an interaction by ID.

        Parameters
        ----------
        interaction_ID : str
            ID of the interaction

        Returns
        -------
        Interaction
            The interaction with the given ID

        Raises
        ------
        ValueError
            If interaction with given ID does not exist
        """
        if interaction_ID in self._self_interaction_dict:
            return self._self_interaction_dict[interaction_ID]
        elif interaction_ID in self._mutual_interaction_dict:
            return self._mutual_interaction_dict[interaction_ID]
        else:
            raise ValueError("Interaction with ID {} does not exist. Existing interactions: {} {}".format(interaction_ID, 
                self._self_interaction_dict.keys(), self._mutual_interaction_dict.keys()))
```

--------------------------------

### Get First Shell Rolling Functions

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/lattice.html

Provides a list of functions to move a site to its nearest neighbors in a simple cubic lattice.

```python
def _1st_shell_roller(self):
        roller = [
            lambda x: jnp.roll(x, 1, axis=0),
            lambda x: jnp.roll(x, 1, axis=1),
            lambda x: jnp.roll(x, 1, axis=2)
        ]
        return roller
```

--------------------------------

### Initialize Langevin Integrator

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/integrator/llg.html

Initializes the LLG Langevin integrator with time step, temperature, damping constant, and gyromagnetic ratio.

```python
class LLLangevinIntegrator(LLIntegrator):
    """
    [For testing purposes: naive geometric integrator.]
    Stochastic Landau-Lifshitz equation of motion.
    
    See Eriksson, Olle, et al. Atomistic spin dynamics: foundations and applications. Oxford university press, 2017, Sec.7.4.5 for details.
    
    The equation of motion is:
    dM/dt = -gamma M x (B + b) - (gamma * alpha / ||M||) * M x (M x (B + b))
    
    b is the stochastic force <b_i_alpha(t) b_j_beta(s)> = 2 * D * delta_ij * delta_alpha_beta * delta(t-s)
    A steady Boltzmann state requires D= alpha/(1+alpha^2) * kbT / gamma / ||m||

    Parameters
    ----------
    dt : float
        Time step size
    temp : float
        Temperature
    alpha : float
        Gilbert damping constant
    gamma : float, optional
        Gyromagnetic ratio. If None, uses electron gyromagnetic ratio
    """
 

    def __init__(self, dt, temp, alpha, gamma=None):
        super().__init__(dt, alpha, gamma)
        self.kbT = Constants.kb * temp
        self.D_base = self.alpha/(1+self.alpha**2) * self.kbT / self.gamma
```

--------------------------------

### Get Second Shell Rolling Functions

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/lattice.html

Provides a list of functions to move a site to its second nearest neighbors in a simple cubic lattice.

```python
def _2nd_shell_roller(self):
        roller = [
            lambda x: jnp.roll(x, (1, 1), axis=(0,1)),
            lambda x: jnp.roll(x, (-1,1), axis=(0,1)),
```

--------------------------------

### Set OpenFerro to Run on CPU

Source: https://openferro.readthedocs.io/en/latest/_sources/quickstart.rst.txt

Specify the JAX_PLATFORMS environment variable to ensure OpenFerro runs on CPU. JAX automatically determines the number of CPU cores to utilize.

```bash
export JAX_PLATFORMS=cpu
```

--------------------------------

### Get Inhomogeneous Strain Dipole Interaction Function

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/engine/ferroelectric.html

Returns the inhomogeneous strain dipole interaction function. This function can be JIT compiled for performance.

```python
def get_inhomo_strain_dipole_interaction(enable_jit=True):
    """
    Returns the inhomogeneous strain dipole interaction function.

    Parameters
    ----------
    enable_jit : bool, optional
        Whether to enable JIT compilation, by default True

    Returns
    -------
    function
        The interaction function
    """
    get_local_strain = jit(LocalStrain3D.get_local_strain) if enable_jit else LocalStrain3D.get_local_strain
    def inhomo_strain_dipole_interaction(local_displacement, dipole_field, parameters):
        """
        Returns the inhomogeneous strain dipole interaction energy.


```

--------------------------------

### Initialize Device Mesh for Multi-GPU Parallelism

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/parallelism.html

Initializes a JAX Mesh for distributing computations across multiple devices (e.g., GPUs) on a single host. Automatically determines mesh dimensions if not specified, or validates user-provided dimensions. Raises ValueError if only one device is available or if the specified dimensions do not match the number of devices.

```python
import logging
import numpy as np
import jax
from jax.sharding import Mesh, PartitionSpec, NamedSharding



class DeviceMesh:


    def __init__(self, devices=None, num_rows=None, num_cols=None):
        """
        Initialize the single-host multi-device parallelism. Get the mesh of the devices.

        Parameters
        ----------
        devices : array-like, optional
            List of devices to use. If None, uses all available devices
        num_rows : int, optional
            Number of rows in device mesh. If None, automatically determined
        num_cols : int, optional
            Number of columns in device mesh. If None, automatically determined

        Raises
        ------
        ValueError
            If only one device is available
            If num_rows * num_cols does not match number of devices
        """
        if devices is None:
            devices = np.array(jax.devices())
        else:
            devices = np.array(devices)
        num_devices = len(devices)
        if num_devices == 1:
            num_rows = 1
            num_cols = 1
            raise ValueError("Only one device is available. No parallelism is applied.")
        if num_rows is None or num_cols is None:
            for i in range(int(np.sqrt(num_devices)), 0, -1):
                if num_devices % i == 0:
                    num_rows = i
                    num_cols = num_devices // i
                    break
        else:
            num_rows = int(num_rows)
            num_cols = int(num_cols)
            if num_rows * num_cols != num_devices:
                raise ValueError("The number of devices does not match the configuration.")
        logging.info('The number of devices is {}'.format(num_devices))
        logging.info('The configuration of the devices is ({} x {})'.format(num_rows, num_cols))
        devices = devices.reshape(num_rows, num_cols)
        # Create a Mesh object to distribute a value across devices:
        self.mesh = Mesh(devices=devices, axis_names=('x', 'y'))

```

--------------------------------

### Run Structural Minimization in OpenFerro

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/simulation.html

Executes structural relaxation, either with a variable or fixed cell. Requires pressure specification for variable cell minimization.

```python
## structural relaxation
self.initialize_reporters()
for i in range(self.max_iter):
    self._step(variable_cell)
    self.step_reporters()
    converged = []
    for field in self.all_fields:
        if jnp.max(jnp.abs(field.get_force())) < self.tol:
            converged.append(True)
        else:
            converged.append(False)
    if all(converged):
        break
```

--------------------------------

### interaction_base Class

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/interaction.html

The base class for all interactions. It provides methods to set and get interaction parameters, and to set the energy and force calculation engines.

```APIDOC
## interaction_base

### Description
The base class to specify the interaction between fields.

### Methods
#### `__init__(self, parameters=None)`
Initializes the interaction base class with optional parameters.

#### `set_parameters(self, parameters)`
Sets the parameters for the interaction. Parameters are converted to a JAX numpy array.

- **Parameters**: `parameters` (array_like) - The parameters of the interaction.

- **Raises**: `ValueError` - If parameters is not a numpy array, list, or jax array.

#### `get_parameters(self)`
Gets the parameters of the interaction.

- **Returns**: `jax.numpy.ndarray` - The parameters of the interaction.

#### `set_energy_engine(self, energy_engine, enable_jit=True)`
Sets the energy engine for the interaction. The energy engine is a callable that takes field values and returns energy. JIT compilation can be enabled.

- **Parameters**:
  - `energy_engine` (callable) - The energy engine function.
  - `enable_jit` (bool, optional) - Whether to enable JIT compilation. Defaults to True.

#### `calc_energy(self)`
Placeholder method to calculate the energy of the interaction. Should be implemented by subclasses.

#### `calc_force(self)`
Placeholder method to calculate the force of the interaction. Should be implemented by subclasses.
```

--------------------------------

### Initialize Simulation

Source: https://openferro.readthedocs.io/en/latest/api.html

The base Simulation class requires a System object upon initialization. It manages time evolution and reporters.

```python
sim = Simulation(system)
```

--------------------------------

### Base Interaction Class

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/interaction.html

The base class for specifying interactions between fields. It includes methods to set and get interaction parameters, and to set the energy and force engines.

```python
class interaction_base:
    """
    The base class to specify the interaction between fields.
    """


    def __init__(self, parameters=None):
        self.parameters = parameters
        self.energy_engine = None
        self.force_engine = None




    def set_parameters(self, parameters):
        """
        Set the parameters of the interaction.

        Parameters
        ----------
        parameters : array_like
            The parameters of the interaction

        Raises
        ------
        ValueError
            If parameters is not a numpy array, list, or jax array
        """
        ## turning list or numpy array to jax array
        if isinstance(parameters, jnp.ndarray):
            paras = parameters
        elif isinstance(parameters, np.ndarray) or isinstance(parameters, list):
            paras = jnp.array(parameters)
        else:
            raise ValueError("Parameters must be a numpy array, a list, or a jax array")
        self.parameters = paras




    def get_parameters(self):
        """
        Get the parameters of the interaction.

        Returns
        -------
        jax.numpy.ndarray
            The parameters of the interaction
        """
        return self.parameters




    def set_energy_engine(self, energy_engine, enable_jit=True):
        """
        Set the energy engine of the interaction.

        Parameters
        ----------
        energy_engine : callable
            The energy engine of the interaction. It should take the values of the fields as input and return the energy as output.
        enable_jit : bool, optional
            Whether to enable jit for the energy engine, by default True
        """
        if enable_jit:
            self.energy_engine = jit(energy_engine)
        else:
            self.energy_engine = energy_engine




    def calc_energy(self):
        pass




    def calc_force(self):
        pass
```

--------------------------------

### Initialize Thermo_Reporter

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/reporter.html

Initializes the thermodynamic reporter, setting up the output file, logging interval, and flags for reporting various quantities like strain, stress, volume, energy, and temperature. It also writes the header information to the log file.

```python
def __init__(self, file='thermo.log', log_interval=100, global_strain=False, excess_stress=False, volume=True, potential_energy=False, kinetic_energy=False, temperature=False):
        self.file = file
        self.counter = -1
        self.log_interval = log_interval
        self.report_global_strain = global_strain
        self.report_excess_stress = excess_stress
        self.report_volume = volume
        self.report_potential_energy = potential_energy
        self.report_kinetic_energy = kinetic_energy
        self.report_temperature = temperature
        self.item_list = []
```

```python
def initialize(self, system):
        """
        Initialize the reporter.

        Parameters
        ----------
        system : System
            The physical system to report on
        """
        ## make the directory if not exists
        all_fields = system.get_all_fields()
        self.item_list = ['Step']
        if self.report_global_strain:
            self.item_list.append('Strain-1')
            self.item_list.append('Strain-2')
            self.item_list.append('Strain-3')
            self.item_list.append('Strain-4')
            self.item_list.append('Strain-5')
            self.item_list.append('Strain-6')
        if self.report_excess_stress:
            self.item_list.append('ex_stress-1')
            self.item_list.append('ex_stress-2')
            self.item_list.append('ex_stress-3')
            self.item_list.append('ex_stress-4')
            self.item_list.append('ex_stress-5')
            self.item_list.append('ex_stress-6')
        if self.report_volume:
            self.item_list.append('Volume')
        if self.report_potential_energy:
            self.item_list.append('Total_Potential_Energy')
        if self.report_kinetic_energy:
            for field in all_fields:
                self.item_list.append('Kinetic-{}'.format(field.ID))
        if self.report_temperature:
            for field in all_fields:
                self.item_list.append('Temperature-{}'.format(field.ID))
        ## write the header: system lattice constants, lattice type, lattice size, etc.
        with open(self.file, 'a') as f:
            ## write the start time
            f.write(f"# Log file for OpenFerro Simulation started at: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}\n")
            f.write(f"# Lattice Type: {system.lattice.name}, Lattice Size: {system.lattice.size}\n")
            f.write(f"# Lattice Vectors: {system.lattice.a1}, {system.lattice.a2}, {system.lattice.a3}\n")
            f.write(f"# Current Time Step: {self.counter+1}\n")
            ## write the report items
            f.write("# ")
            f.write(", ".join(self.item_list))
            f.write("\n")
```

--------------------------------

### System Class Methods

Source: https://openferro.readthedocs.io/en/latest/api.html

Provides an overview of the methods available for the System class, used for managing and simulating physical systems.

```APIDOC
## System Class

### Description
Represents a physical system in OpenFerro, managing fields, interactions, and energy calculations.

### Methods
- `System.lattice`
- `System.__init__()`
- `System.get_field_by_ID(id)`
- `System.get_all_fields()`
- `System.get_all_SO3_fields()`
- `System.get_all_non_SO3_fields()`
- `System.move_fields_to_multi_devs()`
- `System.add_field(field)`
- `System.add_global_strain(strain)`
- `System.get_interaction_by_ID(id)`
- `System._add_interaction_sanity_check()`
- `System.add_dipole_dipole_interaction(...)`
- `System.add_dipole_onsite_interaction(...)`
- `System.add_dipole_interaction_1st_shell(...)`
- `System.add_dipole_interaction_2nd_shell(...)`
- `System.add_dipole_interaction_3rd_shell(...)`
- `System.add_dipole_efield_interaction(...)`
- `System.add_homo_elastic_interaction(...)`
- `System.add_inhomo_elastic_interaction(...)`
- `System.add_pressure(...)`
- `System.add_homo_strain_dipole_interaction(...)`
- `System.add_inhomo_strain_dipole_interaction(...)`
- `System.add_cubic_anisotropy_interaction(...)`
- `System._add_isotropic_exchange_interaction_by_rollers(...)`
- `System.add_isotropic_exchange_interaction_1st_shell(...)`
- `System.add_isotropic_exchange_interaction_2nd_shell(...)`
- `System.add_isotropic_exchange_interaction_3rd_shell(...)`
- `System.add_isotropic_exchange_interaction_4th_shell(...)`
- `System.add_self_interaction(...)`
- `System.add_mutual_interaction(...)`
- `System.add_triple_interaction(...)`
- `System.calc_energy_by_ID(id)`
- `System.calc_force_by_ID(id)`
- `System.calc_total_self_energy()`
- `System.calc_total_mutual_interaction()`
- `System.calc_total_triple_interaction()`
- `System.calc_total_potential_energy()`
- `System.calc_total_kinetic_energy()`
- `System.calc_temp_by_ID(id)`
- `System.calc_excess_stress()`
- `System.update_force_from_self_interaction()`
- `System.update_force_from_mutual_interaction()`
- `System.update_force_from_triple_interaction()`
- `System.update_force()`

### Properties
- `System.interaction_dict`
```

--------------------------------

### Get 3rd Shell Rolling Functions for SimpleSquare2D

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/lattice.html

Returns rolling functions to move a site to all sites in the third shell of a SimpleSquare2D lattice. Used for calculating interactions.

```python
    def _3rd_shell_roller(self):
        """
        Return a list of rolling functions for moving a site to all sites in the third shell of the lattice.

        Returns
        -------
        list
            List of rolling functions. The number of rolling functions is half the shell coordination number.
            Used to calculate interactions between a site and its neighbours.
        """
        roller = [lambda x: jnp.roll(x, 2, axis=i) for i in range(2)]
        return roller
```

--------------------------------

### Get 2nd Shell Rolling Functions for SimpleSquare2D

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/lattice.html

Returns rolling functions to move a site to all sites in the second shell of a SimpleSquare2D lattice. Used for calculating interactions.

```python
    def _2nd_shell_roller(self):
        """
        Return a list of rolling functions for moving a site to all sites in the second shell of the lattice.

        Returns
        -------
        list
            List of rolling functions. The number of rolling functions is half the shell coordination number.
            Used to calculate interactions between a site and its neighbours.
        """
        roller = [
            lambda x: jnp.roll(x, (1,1), axis=(0,1)),
            lambda x: jnp.roll(x, (-1,1), axis=(0,1)),
        ]
        return roller
```

--------------------------------

### Initialize Device Mesh for Parallelism

Source: https://openferro.readthedocs.io/en/latest/api.html

The DeviceMesh class initializes single-host multi-device parallelism by creating a mesh of available devices. It allows specifying the devices, number of rows, and number of columns for the mesh.

```python
class DeviceMesh:
    """Initialize the single-host multi-device parallelism. Get the mesh of the devices."""
    def __init__(_devices =None_, _num_rows =None_, _num_cols =None_):
        # Parameters:
        #   devices (array-like, optional): List of devices to use. If None, uses all available devices
        #   num_rows (int, optional): Number of rows in device mesh. If None, automatically determined
        #   num_cols (int, optional): Number of columns in device mesh. If None, automatically determined
        # Raises:
        #   ValueError: If only one device is available If num_rows * num_cols does not match number of devices
```

--------------------------------

### Get 1st Shell Rolling Functions for SimpleSquare2D

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/lattice.html

Returns rolling functions to move a site to all sites in the first shell of a SimpleSquare2D lattice. Used for calculating interactions.

```python
    def _1st_shell_roller(self):
        """
        Return a list of rolling functions for moving a site to all sites in the first shell of the lattice.

        Returns
        -------
        list
            List of rolling functions. The number of rolling functions is half the shell coordination number.
            Used to calculate interactions between a site and its neighbours.
        """
        roller = [lambda x: jnp.roll(x, 1, axis=i) for i in range(2)]
        return roller
```

--------------------------------

### LLSIBIntegrator

Source: https://openferro.readthedocs.io/en/latest/api.html

Initializes the LLSIBIntegrator for simulating Landau-Lifshitz equations.

```APIDOC
## LLSIBIntegrator

### Description
Initializes the Landau-Lifshitz equation of motion using the SIB scheme.

### Method
__init__

### Parameters
* **dt** (float) – Time step size
* **alpha** (float) – Gilbert damping constant
* **gamma** (float, optional) – Gyromagnetic ratio. Defaults to None (uses electron gyromagnetic ratio).
* **max_iter** (int, optional) – Maximum number of iterations for fixed-point iterations. Defaults to 10.
* **tol** (float, optional) – Tolerance for convergence. Convergence is declared if Average[|M_new - M_old|/Ms averaged over lattice] < tol. Defaults to 1e-06.
```

--------------------------------

### SimulationNVE.__init__

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/simulation.html

Initializes the SimulationNVE class for NVE (microcanonical) molecular dynamics simulations.

```APIDOC
## SimulationNVE.__init__

### Description
Initializes the SimulationNVE class for NVE (microcanonical) molecular dynamics simulations.

### Method
`__init__(system)`

### Parameters
- `system` (System): The physical system to simulate.
```

--------------------------------

### Initialize an SO(3) Field

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/field.html

Initializes an SO(3) field, representing a 3D vector with a fixed magnitude and flexible orientation. Examples include rigid spins or molecular orientations.

```python
field_so3 = FieldSO3(lattice, ID='my_so3_field', unit='T')

```

--------------------------------

### Initializing Field Reporter

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/reporter.html

Initializes the Field_Reporter by setting up file names, logging interval, and writing header information to the average log file. This includes lattice details and report items.

```python
    def initialize(self, system):
        """
        Initialize the reporter.

        Parameters
        ----------
        system : System
            The physical system to report on
        """
        field = system.get_field_by_ID(self.field_ID)  ## check if the field exists
        dim = field
        if self.report_field_average:
            self.item_list = ['Step', 'Average(Vector)']
            with open(self.file_avg_name, 'a') as f:
                ## write the start time
                f.write(f"# Log file for the average of field [{self.field_ID}] started at: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}\n")
                f.write(f"# Lattice Type: {system.lattice.name}, Lattice Size: {system.lattice.size}\n")
                f.write(f"# Lattice Vectors: {system.lattice.a1}, {system.lattice.a2}, {system.lattice.a3}\n")
                f.write(f"# Current Time Step: {self.counter+1}\n")
                ## write the report items
                f.write("# ")
                f.write(", ".join(self.item_list))
                f.write("\n")
```

--------------------------------

### Get Excess Stress from GlobalStrain

Source: https://openferro.readthedocs.io/en/latest/_modules/openferro/field.html

Calculates the excess stress from the GlobalStrain field by dividing the force by the lattice's reference volume and converting to bars. This is useful for stress analysis.

```python
    def get_excess_stress(self):
        return self.get_force() / self.lattice.ref_volume / Constants.bar 

```

--------------------------------

### Set JAX to use CPU

Source: https://openferro.readthedocs.io/en/latest/quickstart.html

Specify the JAX_PLATFORMS environment variable to ensure OpenFerro runs on CPU. JAX automatically manages CPU core utilization.

```bash
export JAX_PLATFORMS=cpu

```