Source code for goo.molecule

from typing import Optional, Tuple
from typing_extensions import override
import numpy as np
import math

from scipy.ndimage import laplace
from scipy.integrate import solve_ivp
from scipy.spatial import KDTree

from goo.utils import *


[docs] class Molecule: """A molecule involved in the diffusion system. Args: name (str): The name of the molecule. conc (float): The initial concentration of the molecule. D (float): The diffusion rate of the molecule. gradient (str, optional): The gradient of the molecule. Defaults to None. """ def __init__( self, name: str, conc: float, D: float, gradient: Optional[str] = None ): self._name = name self._D = D self._conc = conc self._gradient = gradient def __repr__(self): str = (f"Molecule(concentration={self.conc}" f"diffusion_rate={self.D})") return str @property def name(self) -> str: """Name of the molecule. Also defines the name of related forces and collections of effectors. """ return self._name @name.setter def name(self, name: str): self._name = name @property def conc(self) -> float: """The concentration of the molecule.""" return self._conc @conc.setter def conc(self, conc: float): self._conc = conc @property def D(self) -> float: """The diffusion rate of the molecule.""" return self._D @D.setter def D(self, D: float): self._D = D @property def gradient(self) -> str: """The gradient of the molecule.""" return self._gradient @gradient.setter def gradient(self, gradient: str): self._gradient = gradient
[docs] class DiffusionSystem: """A diffusion system that simulates the diffusion of molecules in a 3D grid. It also handles the secretion and sensing of molecular signals by cells. Args: molecules (list[Molecule]): The list of molecules in the system. grid_size (Tuple[int, int, int], optional): The size of the 3D grid. Defaults to (25, 25, 25). grid_center (Tuple[int, int, int], optional): The center of the 3D grid. Defaults to (0, 0, 0). time_step (float, optional): The time step of the simulation. Defaults to 0.1. total_time (int, optional): The total time of the simulation. Defaults to 10. element_size (Tuple[float, float, float], optional): The size of each element in the grid. Defaults to (1.0, 1.0, 1.0). """ def __init__( self, molecules: list[Molecule], grid_size: Tuple[int, int, int] = (50, 50, 50), grid_center: Tuple[int, int, int] = (0, 0, 0), time_step: float = 0.1, total_time: int = 1, element_size=(0.5, 0.5, 0.5) ) -> None: self._molecules = molecules self._grid_size = grid_size self._time_step = time_step self._total_time = total_time self._grid_center = grid_center self._grid_concentrations = None self._kd_tree = None self._element_size = element_size def _build_kdtree(self): """Initialize the grid and build its corresponding KD-Tree.""" x = np.linspace( -(self._grid_size[0] - 1) / 2, (self._grid_size[0] - 1) / 2, self._grid_size[0] ) * self._element_size[0] y = np.linspace( -(self._grid_size[1] - 1) / 2, (self._grid_size[1] - 1) / 2, self._grid_size[1] ) * self._element_size[1] z = np.linspace( -(self._grid_size[2] - 1) / 2, (self._grid_size[2] - 1) / 2, self._grid_size[2] ) * self._element_size[2] grid_points = np.array(np.meshgrid(x, y, z)).T.reshape(-1, 3) grid_points += np.array(self._grid_center) self._kd_tree = KDTree(grid_points) # initialize the grid to store concentrations for each molecule self._grid_concentrations = np.zeros((len(self._molecules), *self._grid_size)) for idx, mol in enumerate(self._molecules): match mol._gradient: case None: continue case "constant": self._grid_concentrations[idx] = mol._conc case "random": variation = 0.1 # 10% variation noise = np.random.normal(0, variation * mol._conc, size=self._grid_size) rand_conc = mol._conc + noise # conc is always positive rand_conc = np.clip(rand_conc, 0, None) self._grid_concentrations[idx] = rand_conc case "center": center_index = self._get_nearest_grid_index(self._grid_center) self._grid_concentrations[idx][center_index] += mol._conc case "linear": x_grad = np.linspace( 0, mol._conc, self._grid_size[0] ).reshape(-1, 1, 1) self._grid_concentrations[idx] = np.tile( x_grad, (1, self._grid_size[1], self._grid_size[2]) ) # add 8 empty points at the corners of the grid in Blender corner_offsets = [ (-1, -1, -1), (-1, -1, 1), (-1, 1, -1), (-1, 1, 1), (1, -1, -1), (1, -1, 1), (1, 1, -1), (1, 1, 1) ] for offset in corner_offsets: corner_position = ( self._grid_center[0] + offset[0] * (self._grid_size[0] - 1) / 2 * self._element_size[0], self._grid_center[1] + offset[1] * (self._grid_size[1] - 1) / 2 * self._element_size[1], self._grid_center[2] + offset[2] * (self._grid_size[2] - 1) / 2 * self._element_size[2] ) bpy.ops.object.empty_add(type='PLAIN_AXES', location=corner_position) return self._kd_tree def _get_nearest_grid_index(self, point): """Get the nearest grid index for the given point.""" return self._kd_tree.query(point)[1] def _get_grid_position(self, index): """Get the coordinates of the given index in the grid.""" return np.unravel_index(index, self._grid_size)
[docs] def update_concentration(self, mol_idx, index, value): """Update the concentration value of a molecule at a given voxel in the grid.""" # Convert flat index to 3D index z_index, y_index, x_index = np.unravel_index(index, self._grid_size) # z_index, y_index, x_index = self._get_grid_position(index) self._grid_concentrations[mol_idx, z_index, y_index, x_index] += value return self._grid_concentrations[mol_idx, z_index, y_index, x_index]
[docs] def get_concentration(self, mol_idx, index): """Get the concentration value of a molecule at a given voxel in the grid.""" z_index, y_index, x_index = np.unravel_index(index, self._grid_size) # z_index, y_index, x_index = self._get_grid_position(index) return self._grid_concentrations[mol_idx, z_index, y_index, x_index]
[docs] def diffuse(self, mol_idx): """Update the concentration of molecules based on diffusion.""" # get grid grid = self._grid_concentrations conc = grid[mol_idx] laplacian = laplace(conc, mode='wrap') conc += self._time_step * self._molecules[mol_idx]._D * laplacian conc = np.clip(conc, 0, None) grid[mol_idx] = conc
[docs] def simulate(self): """Run the diffusion simulation over the total time.""" num_steps = int(self._total_time / self._time_step) for _ in range(num_steps): self.diffuse()
@property def molecules(self) -> list[Molecule]: """The list of molecules in the system.""" return self._molecules @molecules.setter def molecules(self, molecules: list[Molecule]): self._molecules = molecules @property def gradient(self) -> list[(Molecule, str)]: """The gradient of the molecules in the system.""" return self.gradient @gradient.setter def gradient(self, gradient: list[(Molecule, str)]): self.gradient = gradient @property def grid_size(self) -> tuple: """The size of the 3D grid.""" return self._grid_size @grid_size.setter def grid_size(self, grid_size: tuple): self._grid_size = grid_size @property def time_step(self) -> float: """The time step of the simulation.""" return self._time_step @time_step.setter def time_step(self, time_step: float): self._time_step = time_step @property def total_time(self) -> int: """The total time of the simulation.""" return self._total_time @total_time.setter def total_time(self, total_time: int): self._total_time = total_time