Skip to content
Snippets Groups Projects
Select Git revision
  • 1ceda29b67ac83ec27277f2437b3506f56a0ccd2
  • main default protected
  • nicoa96-main-patch-39695
  • nicoa96-main-patch-29647
  • jt/bhtfin
  • jt/bhtnew
  • revert-adeaff5e
  • nicoa96-main-patch-73354
  • nicoa96-main-patch-47348
  • jt/bht
  • jima1
  • jima
  • na/bhtalgorithmus
  • yuhe
  • jn/nasa-data
  • Nicola
  • kuba
17 results

system.py

Blame
  • nguyed99's avatar
    nguyed99 authored
    43f2e8df
    History
    Code owners
    Assign users and groups as approvers for specific file changes. Learn more.
    system.py 1.11 KiB
    """
    This module contains base class for a n-body gravitational system.
    """
    
    from collections.abc import Callable
    from dataclasses import dataclass
    
    import numpy as np
    
    
    @dataclass
    class GravitationalSystem:
        r0: np.ndarray[float]  # shape = (N,d)
        v0: np.ndarray[float]  # shape = (N,d)
        m: np.ndarray[float]  # shape = N
        t: np.ndarray[float]  # shape = n_time_steps
        force: Callable
        solver: Callable
    
        def __post_init__(self):
            "Checking dimensions of inputs"
            assert self.r0.shape == self.v0.shape
            assert self.m.shape[0] == self.r0.shape[0]
    
            solvers = ['verlet', 'bht']
            assert self.solver.__name__ in solvers
    
        def simulation(self):
            """
            Using integrator to compute trajectory in phase space
            """
            p = np.zeros((len(self.t), *self.v0.shape))
            q = np.zeros((len(self.t), *self.r0.shape))
            q[0] = self.r0
            p[0] = self.m * self.v0
    
            for i in range(1, len(self.t)):
                dt = self.t[i] - self.t[i - 1]
                p[i], q[i] = self.solver(self.force, q[i - 1], p[i - 1], self.m, dt)
    
            return self.t, p, q