GravitySimulatorAPI

Gravity simulator API

Source code in grav_sim/api.py

class GravitySimulatorAPI:
    """Gravity simulator API"""

    def __init__(self, c_lib_path: Optional[str] = None) -> None:
        """
        Initialize gravity simulator API

        Parameters
        ----------
        c_lib_path : str, optional
            Path to C library, by default None
        """
        if c_lib_path is not None:
            self.c_lib = utils.load_c_lib(Path(c_lib_path))
        else:
            self.c_lib = utils.load_c_lib()
        utils.initialize_c_lib(self.c_lib)

        # System
        self.BUILT_IN_SYSTEMS = System.BUILT_IN_SYSTEMS

        # Plotting
        self.SOLAR_SYSTEM_COLORS = plotting.SOLAR_SYSTEM_COLORS
        self.plot_quantity_against_time = plotting.plot_quantity_against_time
        self.plot_2d_trajectory = plotting.plot_2d_trajectory
        self.plot_3d_trajectory = plotting.plot_3d_trajectory

        # Simulator
        self.simulator = Simulator(c_lib=self.c_lib)
        self.DAYS_PER_YEAR = self.simulator.DAYS_PER_YEAR
        self.launch_simulation = self.simulator.launch_simulation
        self.launch_cosmological_simulation = (
            self.simulator.launch_cosmological_simulation
        )

        # Parameters
        self.AVAILABLE_ACCELERATION_METHODS = (
            parameters.AccelerationParam.AVAILABLE_ACCELERATION_METHODS
        )
        self.AVAILABLE_INTEGRATORS = parameters.IntegratorParam.AVAILABLE_INTEGRATORS
        self.FIXED_STEP_SIZE_INTEGRATORS = (
            parameters.IntegratorParam.FIXED_STEP_SIZE_INTEGRATORS
        )
        self.ADAPTIVE_STEP_SIZE_INTEGRATORS = (
            parameters.IntegratorParam.ADAPTIVE_STEP_SIZE_INTEGRATORS
        )
        self.AVAILABLE_OUTPUT_METHODS = parameters.OutputParam.AVAILABLE_OUTPUT_METHODS
        self.AVAILABLE_OUTPUT_DTYPE = parameters.OutputParam.AVAILABLE_OUTPUT_DTYPE

    def get_new_system(self) -> System:
        """Create a gravitational system

        Returns
        -------
        System object
        """
        return System(c_lib=self.c_lib)

    def get_new_cosmological_system(self) -> CosmologicalSystem:
        """Create a cosmological system

        Returns
        -------
        CosmologicalSystem object
        """
        return CosmologicalSystem(c_lib=self.c_lib)

    def load_system(
        self,
        file_path: str | Path,
    ) -> System:
        """Load system from a CSV file

        Parameters
        ----------
        file_path : str
            File path to load the system from
        """
        return System.load_system(self.c_lib, file_path)

    def get_built_in_system(self, system_name: str) -> System:
        """Get a built-in gravitational system

        Parameters
        ----------
        system_name : str
            Name of the built-in system to be loaded.
        """
        return System.get_built_in_system(self.c_lib, system_name)

    @staticmethod
    def get_new_parameters() -> Tuple[
        parameters.AccelerationParam,
        parameters.IntegratorParam,
        parameters.OutputParam,
        parameters.Settings,
    ]:
        """Create new simulation parameters

        Returns
        -------
        Tuple of acceleration, integrator, output, and settings parameters
        """
        acceleration_param = parameters.AccelerationParam()
        integrator_param = parameters.IntegratorParam()
        output_param = parameters.OutputParam()
        settings = parameters.Settings()

        return acceleration_param, integrator_param, output_param, settings

    def days_to_years(self, days: float | np.ndarray) -> float | np.ndarray:
        return days / self.simulator.DAYS_PER_YEAR

    def years_to_days(self, years: float | np.ndarray) -> float | np.ndarray:
        return years * self.simulator.DAYS_PER_YEAR

    @staticmethod
    def read_csv_data(
        output_dir: str | Path,
    ) -> Tuple[float, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """Read CSV snapshots from the output directory,
        assuming number of particles and particle_ids
        stays the same

        Parameters
        ----------
        output_dir : str | Path
            Output directory path

        Returns
        -------
        G : float
            Gravitational constant
        time : np.ndarray
            Simulation time of each snapshot
        dt : np.ndarray
            Time step of each snapshot
        particle_ids : np.ndarray
            1D array of Particle IDs
        sol_state : np.ndarray
            3D array of solution state for each snapshot, with shape
            (num_snapshots, num_particles, 7) being
            [m, x, y, z, vx, vy, vz]
        """
        output_dir = Path(output_dir)
        if not output_dir.is_dir():
            raise FileNotFoundError(f"Output directory not found: {output_dir}")

        snapshot_files = sorted(output_dir.glob("snapshot_*.csv"))
        if len(snapshot_files) == 0:
            raise FileNotFoundError(f"No snapshot files found in: {output_dir}")

        G = -1.0
        time = np.zeros(len(snapshot_files), dtype=np.float64)
        dt = np.zeros(len(snapshot_files), dtype=np.float64)

        for i, snapshot_file in enumerate(snapshot_files):
            # Read the metadata
            with open(snapshot_file, "r") as file:
                read_metadata_num_particles = False
                read_metadata_G = False
                read_metadata_time = False
                read_metadata_dt = False
                for line in file:
                    line = line.strip()

                    if line.startswith("#"):
                        if line.startswith("# num_particles"):
                            if i == 0:
                                num_particles = int(line.split(":")[1].strip())
                            elif num_particles != int(line.split(":")[1].strip()):
                                raise ValueError(
                                    f"Number of particles changed from {num_particles} to {int(line.split(':')[1].strip())}"
                                )
                            read_metadata_num_particles = True
                        elif line.startswith("# G"):
                            G = float(line.split(":")[1].strip())
                            read_metadata_G = True
                        elif line.startswith("# time"):
                            time[i] = float(line.split(":")[1].strip())
                            read_metadata_time = True
                        elif line.startswith("# dt"):
                            dt[i] = float(line.split(":")[1].strip())
                            read_metadata_dt = True

                    if (
                        read_metadata_num_particles
                        and read_metadata_G
                        and read_metadata_time
                        and read_metadata_dt
                    ):
                        break

        # Read the data
        particle_ids = np.zeros(num_particles, dtype=np.int32)
        sol_state = np.zeros((len(snapshot_files), num_particles, 7), dtype=np.float64)
        for i, snapshot_file in enumerate(snapshot_files):
            data = np.genfromtxt(snapshot_file, delimiter=",", skip_header=5)
            if i == 0:
                particle_ids = data[:, 0].astype(np.int32)
                particle_ids = np.sort(particle_ids)
                _, num_duplicates = np.unique(particle_ids, return_counts=True)
                if np.any(num_duplicates > 1):
                    raise ValueError(
                        f"Particle IDs are not unique. Particle IDs: {particle_ids}"
                    )

            snapshot_particle_ids = data[:, 0].astype(np.int32)

            # Sort the data by particle IDs
            sorted_indices = np.argsort(snapshot_particle_ids)
            data = data[sorted_indices]

            # Check if the particle IDs match
            if not np.array_equal(particle_ids, snapshot_particle_ids):
                raise ValueError(
                    f"Particle IDs do not match in snapshot {i + 1}: {snapshot_file}"
                )

            # Store the data
            sol_state[i, :, :] = data[:, 1:]

        return G, time, dt, particle_ids, sol_state

    @staticmethod
    def read_hdf5_data(
        output_dir: str | Path,
    ) -> Tuple[float, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """Read HDF5 snapshots from the output directory,
        assuming number of particles and particle_ids
        stays the same

        Parameters
        ----------
        output_dir : str | Path
            Output directory path

        Returns
        -------
        G : float
            Gravitational constant
        time : np.ndarray
            Simulation time of each snapshot
        dt : np.ndarray
            Time step of each snapshot
        particle_ids : np.ndarray
            1D array of Particle IDs
        sol_state : np.ndarray
            3D array of solution state for each snapshot, with shape
            (num_snapshots, num_particles, 7) being
            [m, x, y, z, vx, vy, vz]
        """
        output_dir = Path(output_dir)
        if not output_dir.is_dir():
            raise FileNotFoundError(f"Output directory not found: {output_dir}")

        snapshot_files = sorted(output_dir.glob("snapshot_*.hdf5"))
        if len(snapshot_files) == 0:
            raise FileNotFoundError(f"No snapshot files found in: {output_dir}")

        G = -1.0
        time = np.zeros(len(snapshot_files), dtype=np.float64)
        dt = np.zeros(len(snapshot_files), dtype=np.float64)

        for i, snapshot_file in enumerate(snapshot_files):
            # Read the metadata
            with h5py.File(snapshot_file, "r") as file:
                num_particles = file["Header"].attrs["NumPart_Total"][0]
                G = file["Header"].attrs["G"][0]
                time[i] = file["Header"].attrs["Time"][0]
                dt[i] = file["Header"].attrs["dt"][0]

        # Read the data
        particle_ids = np.zeros(num_particles, dtype=np.int32)
        sol_state = np.zeros((len(snapshot_files), num_particles, 7), dtype=np.float64)
        for i, snapshot_file in enumerate(snapshot_files):
            with h5py.File(snapshot_file, "r") as file:
                particle_ids = file["PartType0"]["ParticleIDs"][()]
                m = file["PartType0"]["Masses"][()]
                x = file["PartType0"]["Coordinates"][()]
                v = file["PartType0"]["Velocities"][()]

                sol_state[i, :, 0] = m
                sol_state[i, :, 1:4] = x
                sol_state[i, :, 4:7] = v

        return G, time, dt, particle_ids, sol_state

    @staticmethod
    def delete_snapshots(
        output_dir: str | Path,
    ):
        """Delete all snapshots in the output directory

        Parameters
        ----------
        output_dir : str | Path
            Output directory path
        """
        output_dir = Path(output_dir)
        if not output_dir.is_dir():
            raise FileNotFoundError(f"Output directory not found: {output_dir}")

        snapshot_files_csv = sorted(output_dir.glob("snapshot_*.csv"))
        for snapshot_file in snapshot_files_csv:
            snapshot_file.unlink()

        snapshot_files_hdf5 = sorted(output_dir.glob("snapshot_*.hdf5"))
        for snapshot_file in snapshot_files_hdf5:
            snapshot_file.unlink()

    def compute_energy(self, sol_state: np.ndarray, G: float) -> np.ndarray:
        """Compute the total energy of the system

        Parameters
        ----------
        sol_state : np.ndarray
            3D array of solution state for each snapshot, with shape
            (num_snapshots, num_particles, 7) being
            [m, x, y, z, vx, vy, vz]
        G : float
            Gravitational constant

        Returns
        -------
        energy : np.ndarray
            1D array of total energy for each snapshot
        """
        # Check the dimension and shape of sol_state
        if len(sol_state.shape) != 3:
            raise ValueError("sol_state must be a 3D array")

        if sol_state.shape[2] != 7:
            raise ValueError(
                "sol_state must have shape (num_snapshots, num_particles, 7)"
            )

        # Compute the total energy
        energy = np.zeros(sol_state.shape[0], dtype=np.float64)
        self.c_lib.compute_energy_python(
            energy.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
            ctypes.c_double(G),
            sol_state.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
            ctypes.c_int32(sol_state.shape[0]),
            ctypes.c_int32(sol_state.shape[1]),
        )

        return energy

    def compute_linear_momentum(
        self,
        sol_state: np.ndarray,
    ) -> np.ndarray:
        """Compute the total linear_momentum of the system

        Parameters
        ----------
        sol_state : np.ndarray
            3D array of solution state for each snapshot, with shape
            (num_snapshots, num_particles, 7) being
            [m, x, y, z, vx, vy, vz]

        Returns
        -------
        linear_momentum : np.ndarray
            1D array of total linear_momentum for each snapshot
        """
        # Check the dimension and shape of sol_state
        if len(sol_state.shape) != 3:
            raise ValueError("sol_state must be a 3D array")

        if sol_state.shape[2] != 7:
            raise ValueError(
                "sol_state must have shape (num_snapshots, num_particles, 7)"
            )

        # Compute the total energy
        linear_momentum = np.zeros(sol_state.shape[0], dtype=np.float64)
        self.c_lib.compute_linear_momentum_python(
            linear_momentum.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
            sol_state.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
            ctypes.c_int32(sol_state.shape[0]),
            ctypes.c_int32(sol_state.shape[1]),
        )

        return linear_momentum

    def compute_angular_momentum(
        self,
        sol_state: np.ndarray,
    ) -> np.ndarray:
        """Compute the total angular_momentum of the system

        Parameters
        ----------
        sol_state : np.ndarray
            3D array of solution state for each snapshot, with shape
            (num_snapshots, num_particles, 7) being
            [m, x, y, z, vx, vy, vz]

        Returns
        -------
        angular_momentum : np.ndarray
            1D array of total angular_momentum for each snapshot
        """
        # Check the dimension and shape of sol_state
        if len(sol_state.shape) != 3:
            raise ValueError("sol_state must be a 3D array")

        if sol_state.shape[2] != 7:
            raise ValueError(
                "sol_state must have shape (num_snapshots, num_particles, 7)"
            )

        # Compute the total energy
        angular_momentum = np.zeros(sol_state.shape[0], dtype=np.float64)
        self.c_lib.compute_angular_momentum_python(
            angular_momentum.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
            sol_state.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
            ctypes.c_int32(sol_state.shape[0]),
            ctypes.c_int32(sol_state.shape[1]),
        )

        return angular_momentum

    @staticmethod
    def compute_eccentricity(
        G: float,
        sol_state: np.ndarray,
    ) -> np.ndarray:
        """Compute the eccentricity using the sol_state array,
        assuming that the first object is the central object

        Parameters
        ----------
        G : float
            Gravitational constant
        sol_state : np.ndarray
            Solution state of the system

        Returns
        -------
        np.ndarray
            Eccentricity of the system at each time step,
            with shape (num_snapshots, num_particles - 1)

        Notes
        -----
        - The function assumes that the first object is the central object.
        - C library function is not used here since this can be done with
            purely numpy vectorized operations. Nevertheless, we may consider
            implementing this in C library in the future.
        """
        num_snapshots = sol_state.shape[0]
        m_0 = sol_state[0, 0, 0]
        m = sol_state[0, 1:, 0]

        eccentricity = np.zeros(num_snapshots)

        x = sol_state[:, 1:, 1:4].copy() - sol_state[:, 0, 1:4].reshape(-1, 1, 3)
        v = sol_state[:, 1:, 4:7].copy() - sol_state[:, 0, 4:7].reshape(-1, 1, 3)

        denom = G * (m_0 + m)[np.newaxis, :, np.newaxis]
        eccentricity = (
            np.cross(v, np.cross(x, v)) / denom
            - x / np.linalg.norm(x, axis=2)[:, :, np.newaxis]
        )
        eccentricity = np.linalg.norm(eccentricity, axis=2)

        return eccentricity

    @staticmethod
    def compute_inclination(sol_state: np.ndarray) -> np.ndarray:
        """Compute the inclination using the sol_state array,
        assuming that the first object is the central object

        Parameters
        ----------
        sol_state : np.ndarray
            Solution state of the system

        Returns
        -------
        np.ndarray
            Inclination of the system at each time step,
            with shape (num_snapshots, num_particles - 1)

        Notes
        -----
        - The function assumes that the first object is the central object.
        - C library function is not used here since this can be done with
          purely numpy vectorized operations. Nevertheless, we may consider
          implementing this in C library in the future.
        """
        num_snapshots = sol_state.shape[0]

        inclination = np.zeros(num_snapshots)

        x = sol_state[:, 1:, 1:4].copy() - sol_state[:, 0, 1:4].reshape(-1, 1, 3)
        v = sol_state[:, 1:, 4:7].copy() - sol_state[:, 0, 4:7].reshape(-1, 1, 3)

        unit_angular_momentum_vector = (
            np.cross(x, v) / np.linalg.norm(np.cross(x, v), axis=2)[:, :, np.newaxis]
        )
        unit_z = np.array([0, 0, 1])

        inclination = np.arccos(np.sum(unit_angular_momentum_vector * unit_z, axis=2))

        return inclination

    @staticmethod
    def plot_rel_energy_error(
        sol_energy: np.ndarray,
        sol_time: np.ndarray,
        is_log_y: bool = True,
        title: Optional[str] = None,
        xlabel: Optional[str] = "Time",
        ylabel: Optional[str] = "$(E_0 - E(t)) / E_0$",
        save_fig: bool = False,
        save_fig_path: Optional[str | Path] = None,
    ) -> None:
        if sol_energy[0] == 0.0:
            warnings.warn("The initial energy is zero.")
        rel_energy_error = np.abs((sol_energy - sol_energy[0]) / sol_energy[0])
        plotting.plot_quantity_against_time(
            quantity=rel_energy_error,
            sol_time=sol_time,
            title=title,
            xlabel=xlabel,
            ylabel=ylabel,
            is_log_y=is_log_y,
            save_fig=save_fig,
            save_fig_path=save_fig_path,
        )

    @staticmethod
    def plot_rel_linear_momentum_error(
        sol_linear_momentum: np.ndarray,
        sol_time: np.ndarray,
        is_log_y: bool = True,
        title: Optional[str] = None,
        xlabel: Optional[str] = "Time",
        ylabel: Optional[str] = "Relative linear momentum error",
        save_fig: bool = False,
        save_fig_path: Optional[str | Path] = None,
    ) -> None:
        if sol_linear_momentum[0] == 0.0:
            warnings.warn("The initial linear momentum is zero.")
        rel_linear_momentum_error = np.abs(
            (sol_linear_momentum - sol_linear_momentum[0]) / sol_linear_momentum[0]
        )
        plotting.plot_quantity_against_time(
            quantity=rel_linear_momentum_error,
            sol_time=sol_time,
            title=title,
            xlabel=xlabel,
            ylabel=ylabel,
            is_log_y=is_log_y,
            save_fig=save_fig,
            save_fig_path=save_fig_path,
        )

    @staticmethod
    def plot_rel_angular_momentum_error(
        sol_angular_momentum: np.ndarray,
        sol_time: np.ndarray,
        is_log_y: bool = True,
        title: Optional[str] = None,
        xlabel: Optional[str] = "Time",
        ylabel: Optional[str] = "$(L_0 - L(t)) / L_0$",
        save_fig: bool = False,
        save_fig_path: Optional[str | Path] = None,
    ) -> None:
        if sol_angular_momentum[0] == 0.0:
            warnings.warn("The initial angular momentum is zero.")
        angular_momentum_error = np.abs(
            (sol_angular_momentum - sol_angular_momentum[0]) / sol_angular_momentum[0]
        )
        plotting.plot_quantity_against_time(
            quantity=angular_momentum_error,
            sol_time=sol_time,
            title=title,
            xlabel=xlabel,
            ylabel=ylabel,
            is_log_y=is_log_y,
            save_fig=save_fig,
            save_fig_path=save_fig_path,
        )

__init__(c_lib_path=None)

Initialize gravity simulator API

Parameters:

Name	Type	Description	Default
c_lib_path	str	Path to C library, by default None	None

Source code in grav_sim/api.py

def __init__(self, c_lib_path: Optional[str] = None) -> None:
    """
    Initialize gravity simulator API

    Parameters
    ----------
    c_lib_path : str, optional
        Path to C library, by default None
    """
    if c_lib_path is not None:
        self.c_lib = utils.load_c_lib(Path(c_lib_path))
    else:
        self.c_lib = utils.load_c_lib()
    utils.initialize_c_lib(self.c_lib)

    # System
    self.BUILT_IN_SYSTEMS = System.BUILT_IN_SYSTEMS

    # Plotting
    self.SOLAR_SYSTEM_COLORS = plotting.SOLAR_SYSTEM_COLORS
    self.plot_quantity_against_time = plotting.plot_quantity_against_time
    self.plot_2d_trajectory = plotting.plot_2d_trajectory
    self.plot_3d_trajectory = plotting.plot_3d_trajectory

    # Simulator
    self.simulator = Simulator(c_lib=self.c_lib)
    self.DAYS_PER_YEAR = self.simulator.DAYS_PER_YEAR
    self.launch_simulation = self.simulator.launch_simulation
    self.launch_cosmological_simulation = (
        self.simulator.launch_cosmological_simulation
    )

    # Parameters
    self.AVAILABLE_ACCELERATION_METHODS = (
        parameters.AccelerationParam.AVAILABLE_ACCELERATION_METHODS
    )
    self.AVAILABLE_INTEGRATORS = parameters.IntegratorParam.AVAILABLE_INTEGRATORS
    self.FIXED_STEP_SIZE_INTEGRATORS = (
        parameters.IntegratorParam.FIXED_STEP_SIZE_INTEGRATORS
    )
    self.ADAPTIVE_STEP_SIZE_INTEGRATORS = (
        parameters.IntegratorParam.ADAPTIVE_STEP_SIZE_INTEGRATORS
    )
    self.AVAILABLE_OUTPUT_METHODS = parameters.OutputParam.AVAILABLE_OUTPUT_METHODS
    self.AVAILABLE_OUTPUT_DTYPE = parameters.OutputParam.AVAILABLE_OUTPUT_DTYPE

compute_angular_momentum(sol_state)

Compute the total angular_momentum of the system

Parameters:

Name	Type	Description	Default
sol_state	ndarray	3D array of solution state for each snapshot, with shape (num_snapshots, num_particles, 7) being [m, x, y, z, vx, vy, vz]	required

Returns:

Name	Type	Description
angular_momentum	ndarray	1D array of total angular_momentum for each snapshot

Source code in grav_sim/api.py

def compute_angular_momentum(
    self,
    sol_state: np.ndarray,
) -> np.ndarray:
    """Compute the total angular_momentum of the system

    Parameters
    ----------
    sol_state : np.ndarray
        3D array of solution state for each snapshot, with shape
        (num_snapshots, num_particles, 7) being
        [m, x, y, z, vx, vy, vz]

    Returns
    -------
    angular_momentum : np.ndarray
        1D array of total angular_momentum for each snapshot
    """
    # Check the dimension and shape of sol_state
    if len(sol_state.shape) != 3:
        raise ValueError("sol_state must be a 3D array")

    if sol_state.shape[2] != 7:
        raise ValueError(
            "sol_state must have shape (num_snapshots, num_particles, 7)"
        )

    # Compute the total energy
    angular_momentum = np.zeros(sol_state.shape[0], dtype=np.float64)
    self.c_lib.compute_angular_momentum_python(
        angular_momentum.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
        sol_state.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
        ctypes.c_int32(sol_state.shape[0]),
        ctypes.c_int32(sol_state.shape[1]),
    )

    return angular_momentum

compute_eccentricity(G, sol_state) staticmethod

Compute the eccentricity using the sol_state array, assuming that the first object is the central object

Parameters:

Name	Type	Description	Default
G	float	Gravitational constant	required
sol_state	ndarray	Solution state of the system	required

Returns:

Type	Description
ndarray	Eccentricity of the system at each time step, with shape (num_snapshots, num_particles - 1)

Notes

• The function assumes that the first object is the central object.
• C library function is not used here since this can be done with purely numpy vectorized operations. Nevertheless, we may consider implementing this in C library in the future.

Source code in grav_sim/api.py

@staticmethod
def compute_eccentricity(
    G: float,
    sol_state: np.ndarray,
) -> np.ndarray:
    """Compute the eccentricity using the sol_state array,
    assuming that the first object is the central object

    Parameters
    ----------
    G : float
        Gravitational constant
    sol_state : np.ndarray
        Solution state of the system

    Returns
    -------
    np.ndarray
        Eccentricity of the system at each time step,
        with shape (num_snapshots, num_particles - 1)

    Notes
    -----
    - The function assumes that the first object is the central object.
    - C library function is not used here since this can be done with
        purely numpy vectorized operations. Nevertheless, we may consider
        implementing this in C library in the future.
    """
    num_snapshots = sol_state.shape[0]
    m_0 = sol_state[0, 0, 0]
    m = sol_state[0, 1:, 0]

    eccentricity = np.zeros(num_snapshots)

    x = sol_state[:, 1:, 1:4].copy() - sol_state[:, 0, 1:4].reshape(-1, 1, 3)
    v = sol_state[:, 1:, 4:7].copy() - sol_state[:, 0, 4:7].reshape(-1, 1, 3)

    denom = G * (m_0 + m)[np.newaxis, :, np.newaxis]
    eccentricity = (
        np.cross(v, np.cross(x, v)) / denom
        - x / np.linalg.norm(x, axis=2)[:, :, np.newaxis]
    )
    eccentricity = np.linalg.norm(eccentricity, axis=2)

    return eccentricity

compute_energy(sol_state, G)

Compute the total energy of the system

Parameters:

Name	Type	Description	Default
sol_state	ndarray	3D array of solution state for each snapshot, with shape (num_snapshots, num_particles, 7) being [m, x, y, z, vx, vy, vz]	required
G	float	Gravitational constant	required

Returns:

Name	Type	Description
energy	ndarray	1D array of total energy for each snapshot

Source code in grav_sim/api.py

def compute_energy(self, sol_state: np.ndarray, G: float) -> np.ndarray:
    """Compute the total energy of the system

    Parameters
    ----------
    sol_state : np.ndarray
        3D array of solution state for each snapshot, with shape
        (num_snapshots, num_particles, 7) being
        [m, x, y, z, vx, vy, vz]
    G : float
        Gravitational constant

    Returns
    -------
    energy : np.ndarray
        1D array of total energy for each snapshot
    """
    # Check the dimension and shape of sol_state
    if len(sol_state.shape) != 3:
        raise ValueError("sol_state must be a 3D array")

    if sol_state.shape[2] != 7:
        raise ValueError(
            "sol_state must have shape (num_snapshots, num_particles, 7)"
        )

    # Compute the total energy
    energy = np.zeros(sol_state.shape[0], dtype=np.float64)
    self.c_lib.compute_energy_python(
        energy.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
        ctypes.c_double(G),
        sol_state.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
        ctypes.c_int32(sol_state.shape[0]),
        ctypes.c_int32(sol_state.shape[1]),
    )

    return energy

compute_inclination(sol_state) staticmethod

Compute the inclination using the sol_state array, assuming that the first object is the central object

Parameters:

Name	Type	Description	Default
sol_state	ndarray	Solution state of the system	required

Returns:

Type	Description
ndarray	Inclination of the system at each time step, with shape (num_snapshots, num_particles - 1)

Notes

• The function assumes that the first object is the central object.
• C library function is not used here since this can be done with purely numpy vectorized operations. Nevertheless, we may consider implementing this in C library in the future.

Source code in grav_sim/api.py

@staticmethod
def compute_inclination(sol_state: np.ndarray) -> np.ndarray:
    """Compute the inclination using the sol_state array,
    assuming that the first object is the central object

    Parameters
    ----------
    sol_state : np.ndarray
        Solution state of the system

    Returns
    -------
    np.ndarray
        Inclination of the system at each time step,
        with shape (num_snapshots, num_particles - 1)

    Notes
    -----
    - The function assumes that the first object is the central object.
    - C library function is not used here since this can be done with
      purely numpy vectorized operations. Nevertheless, we may consider
      implementing this in C library in the future.
    """
    num_snapshots = sol_state.shape[0]

    inclination = np.zeros(num_snapshots)

    x = sol_state[:, 1:, 1:4].copy() - sol_state[:, 0, 1:4].reshape(-1, 1, 3)
    v = sol_state[:, 1:, 4:7].copy() - sol_state[:, 0, 4:7].reshape(-1, 1, 3)

    unit_angular_momentum_vector = (
        np.cross(x, v) / np.linalg.norm(np.cross(x, v), axis=2)[:, :, np.newaxis]
    )
    unit_z = np.array([0, 0, 1])

    inclination = np.arccos(np.sum(unit_angular_momentum_vector * unit_z, axis=2))

    return inclination

compute_linear_momentum(sol_state)

Compute the total linear_momentum of the system

Parameters:

Name	Type	Description	Default
sol_state	ndarray	3D array of solution state for each snapshot, with shape (num_snapshots, num_particles, 7) being [m, x, y, z, vx, vy, vz]	required

Returns:

Name	Type	Description
linear_momentum	ndarray	1D array of total linear_momentum for each snapshot

Source code in grav_sim/api.py

def compute_linear_momentum(
    self,
    sol_state: np.ndarray,
) -> np.ndarray:
    """Compute the total linear_momentum of the system

    Parameters
    ----------
    sol_state : np.ndarray
        3D array of solution state for each snapshot, with shape
        (num_snapshots, num_particles, 7) being
        [m, x, y, z, vx, vy, vz]

    Returns
    -------
    linear_momentum : np.ndarray
        1D array of total linear_momentum for each snapshot
    """
    # Check the dimension and shape of sol_state
    if len(sol_state.shape) != 3:
        raise ValueError("sol_state must be a 3D array")

    if sol_state.shape[2] != 7:
        raise ValueError(
            "sol_state must have shape (num_snapshots, num_particles, 7)"
        )

    # Compute the total energy
    linear_momentum = np.zeros(sol_state.shape[0], dtype=np.float64)
    self.c_lib.compute_linear_momentum_python(
        linear_momentum.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
        sol_state.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
        ctypes.c_int32(sol_state.shape[0]),
        ctypes.c_int32(sol_state.shape[1]),
    )

    return linear_momentum

delete_snapshots(output_dir) staticmethod

Delete all snapshots in the output directory

Parameters:

Name	Type	Description	Default
output_dir	str | Path	Output directory path	required

Source code in grav_sim/api.py

@staticmethod
def delete_snapshots(
    output_dir: str | Path,
):
    """Delete all snapshots in the output directory

    Parameters
    ----------
    output_dir : str | Path
        Output directory path
    """
    output_dir = Path(output_dir)
    if not output_dir.is_dir():
        raise FileNotFoundError(f"Output directory not found: {output_dir}")

    snapshot_files_csv = sorted(output_dir.glob("snapshot_*.csv"))
    for snapshot_file in snapshot_files_csv:
        snapshot_file.unlink()

    snapshot_files_hdf5 = sorted(output_dir.glob("snapshot_*.hdf5"))
    for snapshot_file in snapshot_files_hdf5:
        snapshot_file.unlink()

get_built_in_system(system_name)

Get a built-in gravitational system

Parameters:

Name	Type	Description	Default
system_name	str	Name of the built-in system to be loaded.	required

Source code in grav_sim/api.py

def get_built_in_system(self, system_name: str) -> System:
    """Get a built-in gravitational system

    Parameters
    ----------
    system_name : str
        Name of the built-in system to be loaded.
    """
    return System.get_built_in_system(self.c_lib, system_name)

get_new_cosmological_system()

Create a cosmological system

Returns:

Type	Description
CosmologicalSystem object

Source code in grav_sim/api.py

def get_new_cosmological_system(self) -> CosmologicalSystem:
    """Create a cosmological system

    Returns
    -------
    CosmologicalSystem object
    """
    return CosmologicalSystem(c_lib=self.c_lib)

get_new_parameters() staticmethod

Create new simulation parameters

Returns:

Type	Description
Tuple of acceleration, integrator, output, and settings parameters

Source code in grav_sim/api.py

@staticmethod
def get_new_parameters() -> Tuple[
    parameters.AccelerationParam,
    parameters.IntegratorParam,
    parameters.OutputParam,
    parameters.Settings,
]:
    """Create new simulation parameters

    Returns
    -------
    Tuple of acceleration, integrator, output, and settings parameters
    """
    acceleration_param = parameters.AccelerationParam()
    integrator_param = parameters.IntegratorParam()
    output_param = parameters.OutputParam()
    settings = parameters.Settings()

    return acceleration_param, integrator_param, output_param, settings

get_new_system()

Create a gravitational system

Returns:

Type	Description
System object

Source code in grav_sim/api.py

def get_new_system(self) -> System:
    """Create a gravitational system

    Returns
    -------
    System object
    """
    return System(c_lib=self.c_lib)

load_system(file_path)

Load system from a CSV file

Parameters:

Name	Type	Description	Default
file_path	str	File path to load the system from	required

Source code in grav_sim/api.py

def load_system(
    self,
    file_path: str | Path,
) -> System:
    """Load system from a CSV file

    Parameters
    ----------
    file_path : str
        File path to load the system from
    """
    return System.load_system(self.c_lib, file_path)

read_csv_data(output_dir) staticmethod

Read CSV snapshots from the output directory, assuming number of particles and particle_ids stays the same

Parameters:

Name	Type	Description	Default
output_dir	str | Path	Output directory path	required

Returns:

Name	Type	Description
G	float	Gravitational constant
time	ndarray	Simulation time of each snapshot
dt	ndarray	Time step of each snapshot
particle_ids	ndarray	1D array of Particle IDs
sol_state	ndarray	3D array of solution state for each snapshot, with shape (num_snapshots, num_particles, 7) being [m, x, y, z, vx, vy, vz]

Source code in grav_sim/api.py

@staticmethod
def read_csv_data(
    output_dir: str | Path,
) -> Tuple[float, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Read CSV snapshots from the output directory,
    assuming number of particles and particle_ids
    stays the same

    Parameters
    ----------
    output_dir : str | Path
        Output directory path

    Returns
    -------
    G : float
        Gravitational constant
    time : np.ndarray
        Simulation time of each snapshot
    dt : np.ndarray
        Time step of each snapshot
    particle_ids : np.ndarray
        1D array of Particle IDs
    sol_state : np.ndarray
        3D array of solution state for each snapshot, with shape
        (num_snapshots, num_particles, 7) being
        [m, x, y, z, vx, vy, vz]
    """
    output_dir = Path(output_dir)
    if not output_dir.is_dir():
        raise FileNotFoundError(f"Output directory not found: {output_dir}")

    snapshot_files = sorted(output_dir.glob("snapshot_*.csv"))
    if len(snapshot_files) == 0:
        raise FileNotFoundError(f"No snapshot files found in: {output_dir}")

    G = -1.0
    time = np.zeros(len(snapshot_files), dtype=np.float64)
    dt = np.zeros(len(snapshot_files), dtype=np.float64)

    for i, snapshot_file in enumerate(snapshot_files):
        # Read the metadata
        with open(snapshot_file, "r") as file:
            read_metadata_num_particles = False
            read_metadata_G = False
            read_metadata_time = False
            read_metadata_dt = False
            for line in file:
                line = line.strip()

                if line.startswith("#"):
                    if line.startswith("# num_particles"):
                        if i == 0:
                            num_particles = int(line.split(":")[1].strip())
                        elif num_particles != int(line.split(":")[1].strip()):
                            raise ValueError(
                                f"Number of particles changed from {num_particles} to {int(line.split(':')[1].strip())}"
                            )
                        read_metadata_num_particles = True
                    elif line.startswith("# G"):
                        G = float(line.split(":")[1].strip())
                        read_metadata_G = True
                    elif line.startswith("# time"):
                        time[i] = float(line.split(":")[1].strip())
                        read_metadata_time = True
                    elif line.startswith("# dt"):
                        dt[i] = float(line.split(":")[1].strip())
                        read_metadata_dt = True

                if (
                    read_metadata_num_particles
                    and read_metadata_G
                    and read_metadata_time
                    and read_metadata_dt
                ):
                    break

    # Read the data
    particle_ids = np.zeros(num_particles, dtype=np.int32)
    sol_state = np.zeros((len(snapshot_files), num_particles, 7), dtype=np.float64)
    for i, snapshot_file in enumerate(snapshot_files):
        data = np.genfromtxt(snapshot_file, delimiter=",", skip_header=5)
        if i == 0:
            particle_ids = data[:, 0].astype(np.int32)
            particle_ids = np.sort(particle_ids)
            _, num_duplicates = np.unique(particle_ids, return_counts=True)
            if np.any(num_duplicates > 1):
                raise ValueError(
                    f"Particle IDs are not unique. Particle IDs: {particle_ids}"
                )

        snapshot_particle_ids = data[:, 0].astype(np.int32)

        # Sort the data by particle IDs
        sorted_indices = np.argsort(snapshot_particle_ids)
        data = data[sorted_indices]

        # Check if the particle IDs match
        if not np.array_equal(particle_ids, snapshot_particle_ids):
            raise ValueError(
                f"Particle IDs do not match in snapshot {i + 1}: {snapshot_file}"
            )

        # Store the data
        sol_state[i, :, :] = data[:, 1:]

    return G, time, dt, particle_ids, sol_state

read_hdf5_data(output_dir) staticmethod

Read HDF5 snapshots from the output directory, assuming number of particles and particle_ids stays the same

Parameters:

Name	Type	Description	Default
output_dir	str | Path	Output directory path	required

Returns:

Name	Type	Description
G	float	Gravitational constant
time	ndarray	Simulation time of each snapshot
dt	ndarray	Time step of each snapshot
particle_ids	ndarray	1D array of Particle IDs
sol_state	ndarray	3D array of solution state for each snapshot, with shape (num_snapshots, num_particles, 7) being [m, x, y, z, vx, vy, vz]

Source code in grav_sim/api.py

@staticmethod
def read_hdf5_data(
    output_dir: str | Path,
) -> Tuple[float, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Read HDF5 snapshots from the output directory,
    assuming number of particles and particle_ids
    stays the same

    Parameters
    ----------
    output_dir : str | Path
        Output directory path

    Returns
    -------
    G : float
        Gravitational constant
    time : np.ndarray
        Simulation time of each snapshot
    dt : np.ndarray
        Time step of each snapshot
    particle_ids : np.ndarray
        1D array of Particle IDs
    sol_state : np.ndarray
        3D array of solution state for each snapshot, with shape
        (num_snapshots, num_particles, 7) being
        [m, x, y, z, vx, vy, vz]
    """
    output_dir = Path(output_dir)
    if not output_dir.is_dir():
        raise FileNotFoundError(f"Output directory not found: {output_dir}")

    snapshot_files = sorted(output_dir.glob("snapshot_*.hdf5"))
    if len(snapshot_files) == 0:
        raise FileNotFoundError(f"No snapshot files found in: {output_dir}")

    G = -1.0
    time = np.zeros(len(snapshot_files), dtype=np.float64)
    dt = np.zeros(len(snapshot_files), dtype=np.float64)

    for i, snapshot_file in enumerate(snapshot_files):
        # Read the metadata
        with h5py.File(snapshot_file, "r") as file:
            num_particles = file["Header"].attrs["NumPart_Total"][0]
            G = file["Header"].attrs["G"][0]
            time[i] = file["Header"].attrs["Time"][0]
            dt[i] = file["Header"].attrs["dt"][0]

    # Read the data
    particle_ids = np.zeros(num_particles, dtype=np.int32)
    sol_state = np.zeros((len(snapshot_files), num_particles, 7), dtype=np.float64)
    for i, snapshot_file in enumerate(snapshot_files):
        with h5py.File(snapshot_file, "r") as file:
            particle_ids = file["PartType0"]["ParticleIDs"][()]
            m = file["PartType0"]["Masses"][()]
            x = file["PartType0"]["Coordinates"][()]
            v = file["PartType0"]["Velocities"][()]

            sol_state[i, :, 0] = m
            sol_state[i, :, 1:4] = x
            sol_state[i, :, 4:7] = v

    return G, time, dt, particle_ids, sol_state

Comments