stellar-spice 0.2__tar.gz

stellar_spice-0.2/PKG-INFO

@@ -0,0 +1,32 @@
+Metadata-Version: 2.1
+Name: stellar-spice
+Version: 0.2
+Description-Content-Type: text/markdown
+Requires-Dist: jax~=0.4.24
+Requires-Dist: jaxlib~=0.4.24
+Requires-Dist: numpy~=1.26.4
+Requires-Dist: flax~=0.8.1
+Requires-Dist: joblib~=1.3.2
+Requires-Dist: transformer-payne~=0.5
+Provides-Extra: huggingface
+Requires-Dist: huggingface-hub~=0.21.3; extra == "huggingface"
+Provides-Extra: dev
+Requires-Dist: pytest~=8.0.0; extra == "dev"
+Requires-Dist: pytest-datadir~=1.5.0; extra == "dev"
+Requires-Dist: coverage~=7.4.3; extra == "dev"
+
+![spice logo](docs/img/spice.svg)
+
+A stellar spectra integration library.
+
+I am going to publish the package on PyPI very soon, for the time being it can be installed from GitHub:
+
+```bash
+pip install git+https://github.com/maja-jablonska/spice
+```
+
+Paper is currently in prep.
+
+## Dev log
+
+Notebooks are being updated one by one.

stellar_spice-0.2/README.md

@@ -0,0 +1,15 @@
+![spice logo](docs/img/spice.svg)
+
+A stellar spectra integration library.
+
+I am going to publish the package on PyPI very soon, for the time being it can be installed from GitHub:
+
+```bash
+pip install git+https://github.com/maja-jablonska/spice
+```
+
+Paper is currently in prep.
+
+## Dev log
+
+Notebooks are being updated one by one.

stellar_spice-0.2/pyproject.toml

@@ -0,0 +1,33 @@
+[build-system]
+requires = ["setuptools >= 61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "stellar-spice"
+version = "0.2"
+readme = "README.md"
+dependencies = [
+    "jax~=0.4.24",
+    "jaxlib~=0.4.24",
+    "numpy~=1.26.4",
+    "flax~=0.8.1",
+    "joblib~=1.3.2",
+    "transformer-payne~=0.5"
+]
+
+[project.optional-dependencies]
+huggingface = ["huggingface-hub~=0.21.3"]
+dev = ["pytest~=8.0.0", "pytest-datadir~=1.5.0", "coverage~=7.4.3"]
+
+[tool.setuptools.dynamic]
+dependencies = {file = ["requirements.txt"]}
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+python_classes = ["Test", "Describe"]
+python_functions = ["test_", "it_", "and_", "but_", "they_"]
+python_files = ["test_*.py",]
+testpaths = ["tests",]
+pythonpath = ["src"]

stellar_spice-0.2/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

stellar_spice-0.2/src/spice/geometry/__init__.py

@@ -0,0 +1,2 @@
+from .utils import polygon_area
+from .sutherland_hodgman import clip

stellar_spice-0.2/src/spice/geometry/sutherland_hodgman.py

@@ -0,0 +1,142 @@
+import jax
+import jax.numpy as jnp
+from jax.typing import ArrayLike
+from spice.geometry.utils import inside, last_non_nan, append_to_last_nan, repeat_last, sort_xy
+
+
+@jax.jit
+def x_y_first_line_vertical(p1, p2, q1, q2):
+    x = p1[0]
+    m2 = (q2[1]-q1[1])/(q2[0]-q1[0])
+    b2 = q1[1]-m2*q1[0]
+    return x, m2*x+b2
+
+@jax.jit
+def x_y_second_line_vertical(p1, p2, q1, q2):
+    x = q1[0]
+    m1 = (p2[1]-p1[1])/(p2[0]-p1[0])
+    b1 = p1[1]-m1*p1[0]
+    return x, m1*x+b1
+
+@jax.jit
+def x_y_no_vertical(p1, p2, q1, q2):
+    m1 = (p2[1]-p1[1])/(p2[0]-p1[0])
+    b1 = p1[1]-m1*p1[0]
+
+    m2 = (q2[1]-q1[1])/(q2[0]-q1[0])
+    b2 = q1[1]-m2*q1[0]
+
+    x = (b2-b1)/(m1-m2)
+    return x, m1*x+b1
+
+
+@jax.jit
+def compute_intersection(p1: ArrayLike, p2: ArrayLike, q1: ArrayLike, q2: ArrayLike) -> ArrayLike:
+    """Calculate the intersection of two lines
+
+    Args:
+        p1 (ArrayLike): point 1 of line 1
+        p2 (ArrayLike): point 2 of line 1
+        q1 (ArrayLike): point 1 of line 2
+        q2 (ArrayLike): point 2 of line 2
+
+    Returns:
+        ArrayLike: coordinates of the intersection of two lines
+    """
+    x, y = jax.lax.cond(jnp.isclose(p1[0]-p2[0], 0),
+                        x_y_first_line_vertical,
+                        lambda a, b, c, d: jax.lax.cond(
+                            jnp.isclose(c[0]-d[0], 0),
+                            x_y_second_line_vertical,
+                            x_y_no_vertical,
+                            a, b, c, d),
+                        p1, p2, q1, q2)
+    
+    return jnp.array([[x, y]])
+
+def s_only_edge_start_inside(final_polygon, c_edge_start, c_edge_end, s_edge_start, s_edge_end):
+    intersection = compute_intersection(s_edge_start,
+                                        s_edge_end,
+                                        c_edge_start,
+                                        c_edge_end)
+    intersection = jax.lax.cond(jnp.any(jnp.isnan(intersection)),
+                                lambda: jnp.nan*intersection,
+                                lambda: intersection)
+    return append_to_last_nan(final_polygon, intersection)
+
+def s_only_edge_end_inside(final_polygon, c_edge_start, c_edge_end, s_edge_start, s_edge_end):
+    intersection = compute_intersection(s_edge_start,
+                                        s_edge_end,
+                                        c_edge_start,
+                                        c_edge_end)
+    intersection = jax.lax.cond(jnp.any(jnp.isnan(intersection)),
+                            lambda: jnp.nan*intersection,
+                            lambda: intersection)
+    polygon_with_intersection = append_to_last_nan(final_polygon, intersection)
+    return append_to_last_nan(polygon_with_intersection, s_edge_end)
+
+def s_both_start_and_end_inside(final_polygon, c_edge_start, c_edge_end, s_edge_start, s_edge_end):
+    return append_to_last_nan(final_polygon, s_edge_end)
+
+
+@jax.jit
+def clip(subject_polygon: ArrayLike, clipping_polygon: ArrayLike) -> ArrayLike:
+    """Calculate the intersection of the triangles subject_polygon and clipping_polygon
+
+    Args:
+        subject_polygon (ArrayLike): x and y coordinates (3, 2)
+        clipping_polygon (ArrayLike): x and y coordinates (3, 2)
+
+    Returns:
+        ArrayLike: intersection between two triangles (6, 2)
+    """ 
+    subject_polygon = sort_xy(subject_polygon)
+    clipping_polygon = sort_xy(clipping_polygon)
+    final_polygon = repeat_last(jnp.concatenate([jnp.copy(subject_polygon), jnp.nan*jnp.ones((6-subject_polygon.shape[0], 2))], axis=0))
+
+    def outer_iteration(i, final_polygon):
+        next_polygon = jnp.copy(final_polygon)
+
+        final_polygon = jnp.nan*jnp.ones((6, 2))
+
+        #c_edge_start = clipping_polygon[((i-1)%(clipping_polygon.shape[0])).astype(int)]
+        c_edge_start = jax.lax.cond(i==0,
+                                    lambda: last_non_nan(clipping_polygon),
+                                    lambda: clipping_polygon[i-1])
+        c_edge_end = clipping_polygon[i]
+
+        def inner_iteration(j, final_polygon):
+            s_edge_start = jax.lax.cond(j==0,
+                                        lambda: last_non_nan(next_polygon),
+                                        lambda: next_polygon[j-1])
+            #s_edge_start = next_polygon[((j-1)%(next_polygon.shape[0])).astype(int)]
+            s_edge_end = next_polygon[j]
+
+            return jax.lax.cond(
+                jnp.any(inside(c_edge_start, c_edge_end, s_edge_end)),
+                lambda: jax.lax.cond(
+                    jnp.any(inside(c_edge_start, c_edge_end, s_edge_start)),
+                    lambda: s_both_start_and_end_inside(
+                        final_polygon,
+                        c_edge_start, c_edge_end,
+                        s_edge_start, s_edge_end
+                    ),
+                    lambda: s_only_edge_end_inside(
+                        final_polygon,
+                        c_edge_start, c_edge_end,
+                        s_edge_start, s_edge_end
+                    ),
+                ),
+                lambda: jax.lax.cond(
+                    jnp.any(inside(c_edge_start, c_edge_end, s_edge_start)),
+                    lambda: s_only_edge_start_inside(
+                        final_polygon,
+                        c_edge_start, c_edge_end,
+                        s_edge_start, s_edge_end
+                    ),
+                    lambda: final_polygon
+                )
+            )
+        
+        return jax.lax.fori_loop(0, next_polygon.shape[0], inner_iteration, final_polygon)
+    return repeat_last(jax.lax.fori_loop(0, clipping_polygon.shape[0], outer_iteration, final_polygon))

stellar_spice-0.2/src/spice/geometry/utils.py

@@ -0,0 +1,167 @@
+import jax
+import jax.numpy as jnp
+from jax.typing import ArrayLike
+
+
+def sort_xy(points: ArrayLike) -> ArrayLike:
+    """Sort points clockwise
+
+    Args:
+        points (ArrayLike): x, y coordinates (n, 2)
+
+    Returns:
+        ArrayLike: points with x, y coordinates sorted clockwise
+    """
+    x, y = points[:, 0], points[:, 1]
+    x0 = jnp.mean(x)
+    y0 = jnp.mean(y)
+
+    r = jnp.sqrt((x-x0)**2 + (y-y0)**2)
+
+    angles = jnp.where((y-y0) > 0, jnp.arccos((x-x0)/r), 2*jnp.pi-jnp.arccos((x-x0)/r))
+
+    mask = jnp.argsort(angles)
+
+    return points[mask, :]
+
+
+
+def inside(p1: ArrayLike, p2: ArrayLike, q: ArrayLike) -> bool:
+    """Check if the point is inside the polygon edge
+
+    Args:
+        p1 (ArrayLike): Polygon edge point 1
+        p2 (ArrayLike): Polygon edge point 2
+        q (ArrayLike): Point to be checked
+
+    Returns:
+        bool: Point is inside polygon
+    """
+    return (p2[0]-p1[0])*(q[1]-p1[1])-(p2[1]-p1[1])*(q[0]-p1[0])>=0
+
+
+def last_non_nan_arg(arr: ArrayLike) -> ArrayLike:
+    """Return the last non-nan index of the array
+
+    Args:
+        arr (ArrayLike): Array
+
+    Returns:
+        ArrayLike: Last non-nan index of the array
+    """    
+    return jnp.max(jnp.argwhere(jnp.all(~jnp.isnan(arr), axis=1), size=arr.shape[0], fill_value=0)).astype(int)
+
+
+def last_non_nan(arr: ArrayLike, axis=1) -> ArrayLike:
+    """Return the last non-nan value of the array
+
+    Args:
+        arr (ArrayLike): Array
+
+    Returns:
+        ArrayLike: Last non-nan row of the array
+    """    
+    return arr[jnp.max(jnp.argwhere(jnp.all(~jnp.isnan(arr), axis=axis), size=arr.shape[0], fill_value=0))]
+
+
+def append_to_last_nan(arr: ArrayLike, to_add: ArrayLike) -> ArrayLike:
+    """Append the row to the first nan row
+
+    Args:
+        arr (ArrayLike): Array to add the row to
+        to_add (ArrayLike): The row to be added
+
+    Returns:
+        ArrayLike: Array with an appended row
+    """    
+    next_ind = jnp.min(jnp.argwhere(jnp.all(jnp.isnan(arr), axis=1), size=arr.shape[0], fill_value=arr.shape[0]-1))
+    return arr.at[next_ind].set(to_add.flatten())
+
+
+def append_value_to_last_nan(arr: ArrayLike, to_add: ArrayLike) -> ArrayLike:
+    """Append a single value to the first nan in array
+
+    Args:
+        arr (ArrayLike): Array to add the value to
+        to_add (ArrayLike): The value to be appended
+
+    Returns:
+        ArrayLike: Array with an appended value
+    """    
+    next_ind = jnp.min(jnp.argwhere(jnp.isnan(arr), size=arr.shape[0], fill_value=arr.shape[0]-1))
+    return arr.at[next_ind].set(to_add)
+
+
+def repeat_last(arr: ArrayLike) -> ArrayLike:
+    """Repeat the last non-nan row to fill the whole array
+
+    Args:
+        arr (ArrayLike):
+
+    Returns:
+        ArrayLike: Array with nan values filled with last non-nan values
+    """    
+    return jax.lax.cond(jnp.any(jnp.all(jnp.isnan(arr), axis=1)),
+                        lambda: jax.lax.fori_loop(0, jnp.sum(jnp.all(jnp.isnan(arr), axis=1)),
+                                                  lambda i, _arr: append_to_last_nan(_arr, last_non_nan(_arr)),
+                                                  arr),
+                        lambda: arr)
+
+
+def wrap(arr: ArrayLike) -> ArrayLike:
+    """Copy the first row and concatenate to the first array
+
+    Args:
+        arr (ArrayLike):
+
+    Returns:
+        ArrayLike: Array with the copy of the first row appended
+    """    
+    return jnp.concatenate([arr, arr[0].reshape((-1, arr.shape[1]))])
+
+
+def polygon_area(x: ArrayLike, y: ArrayLike) -> ArrayLike:
+    """Calculate a polygon's surface area from its x and y coordinates
+
+    Args:
+        x (ArrayLike): x coordinates of the polygon
+        y (ArrayLike): y coordinates of the polygon
+
+    Returns:
+        ArrayLike: polygon's surface area
+    """    
+    return 0.5*jnp.abs(jnp.dot(x,jnp.roll(y,1))-jnp.dot(y,jnp.roll(x,1)))
+
+
+polygon_areas = jax.jit(jax.vmap(polygon_area, in_axes=(0, 0)))
+
+def cast_indexes(cast_vertices: ArrayLike) -> ArrayLike:
+    """When casting 3D vectors to 2D, get the indices of non-reduced dimensions
+
+    Args:
+        cast_vertices (ArrayLike): vertices after 3D to 2D casting
+
+    Returns:
+        ArrayLike: non-reduced (non-zero) dimensions
+    """
+    zeroth_ind = jnp.argmax(jnp.all(jnp.isclose(cast_vertices, 0), axis=1))
+
+    return jax.lax.cond(zeroth_ind == 0,
+                        lambda: jnp.array([1, 2]),
+                        lambda: jax.lax.cond(zeroth_ind == 1,
+                                     lambda: jnp.array([0, 2]),
+                                     lambda: jnp.array([0, 1])))
+
+
+def get_cast_areas(face_vertices: ArrayLike) -> ArrayLike:
+    """Get visible, 2D-casted polygon areas
+
+    Args:
+        face_vertices (ArrayLike): polygons' vertices after 3D to 2D casting
+
+    Returns:
+        ArrayLike: areas of 2D-casted polygons
+    """
+    cast_ind = cast_indexes(face_vertices)
+    return polygon_areas(face_vertices[:, :, cast_ind[0]],
+                         face_vertices[:, :, cast_ind[1]])

stellar_spice-0.2/src/spice/models/__init__.py

@@ -0,0 +1,5 @@
+from .mesh_generation import icosphere
+from .mesh_model import IcosphereModel, MeshModel
+from .spots import add_spot, add_spots
+from .binary import Binary
+from .orbit_utils import get_orbit_jax

stellar_spice-0.2/src/spice/models/binary.py

@@ -0,0 +1,123 @@
+import jax.numpy as jnp
+from functools import partial
+import jax
+from jax.typing import ArrayLike
+from typing import NamedTuple, Tuple
+from .mesh_model import MeshModel
+from .mesh_transform import transform, evaluate_body_orbit
+import astropy.units as u
+from .orbit_utils import get_orbit_jax
+from .mesh_view import resolve_occlusion, Grid
+
+
+YEAR_TO_SECONDS = (u.year).to(u.s)
+
+
+class Binary(NamedTuple):
+    body1: MeshModel
+    body2: MeshModel
+
+    # Orbital elements
+    P: float
+    ecc: float
+    T: float
+    i: float
+    omega: float
+    Omega: float
+
+    evaluated_times: ArrayLike
+
+    body1_centers: ArrayLike
+    body2_centers: ArrayLike
+    body1_velocities: ArrayLike
+    body2_velocities: ArrayLike
+
+    @classmethod
+    def from_bodies(cls, body1: MeshModel, body2: MeshModel) -> "Binary":
+          """Construct a Binary object from two mesh models.
+
+          Args:
+              body1 (MeshModel):
+              body2 (MeshModel):
+
+          Returns:
+              Binary: a binary consisting of body1 and body2
+          """
+          return cls(body1, body2, 1., 0., 0., 0., 0., 0., 0., jnp.zeros_like(body1.centers), jnp.zeros_like(body2.centers), jnp.zeros_like(body1.velocities), jnp.zeros_like(body2.velocities))
+
+
+@partial(jax.jit, static_argnums=(7,))
+def add_orbit(binary: Binary, P: float, ecc: float,
+              T: float, i: float, omega: float, Omega: float, orbit_resolution_points: int) -> Binary:
+      """Add orbit information to the binary object
+
+      Args:
+          binary (Binary):
+          P (float): orbit period [years]
+          ecc (float): orbit eccentrity
+          T (float): [defined time units]
+          i (float): inclination [rad]
+          omega (float): []
+          Omega (float): []
+          orbit_resolution_points (int): number of times to resolve the orbit at
+
+      Returns:
+          Binary: object with calculated orbit property values
+      """
+
+      orbit_resolution_times = jnp.linspace(0, P, orbit_resolution_points)
+      orbit = get_orbit_jax(orbit_resolution_times, binary.body1.mass, binary.body2.mass, P,
+                        ecc, T, i, omega, Omega)
+      return binary._replace(P=P, ecc=ecc, T=T, i=i, omega=omega, Omega=Omega,
+                              evaluated_times=orbit_resolution_times,
+                              body1_centers=orbit[2, :, :], body2_centers=orbit[4, :, :],
+                              body1_velocities=orbit[3, :, :], body2_velocities=orbit[5, :, :])
+
+
+@partial(jax.jit, static_argnums=(2,))
+def _evaluate_orbit(binary: Binary, time: ArrayLike, grid: Grid) -> Tuple[MeshModel, MeshModel]:
+      interpolate_orbit = jax.jit(jax.vmap(lambda x: jnp.interp(time, binary.evaluated_times, x, period=binary.P), in_axes=(0,)))
+      body1_center = interpolate_orbit(binary.body1_centers)
+      body2_center = interpolate_orbit(binary.body2_centers)
+      body1_velocity = interpolate_orbit(binary.body1_velocities)
+      body2_velocity = interpolate_orbit(binary.body2_velocities)
+
+      body1 = evaluate_body_orbit(transform(binary.body1, binary.body1.center+body1_center), body1_velocity)
+      body2 = evaluate_body_orbit(transform(binary.body2, binary.body2.center+body2_center), body2_velocity)
+      
+      return jax.lax.cond(jnp.mean(body1.los_z)>jnp.mean(body2.los_z),
+                          lambda: (body1, resolve_occlusion(body2, body1, grid)),
+                          lambda: (resolve_occlusion(body1, body2, grid), body2))
+      
+      
+v_evaluate_orbit = jax.jit(jax.vmap(_evaluate_orbit, in_axes=(None, 0, None)), static_argnums=(2,))
+
+
+def evaluate_orbit(binary: Binary, time: ArrayLike, n_cells: int = 20) -> Tuple[MeshModel, MeshModel]:
+      """_summary_
+
+      Args:
+          binary (Binary): a binary object
+          time (ArrayLike): time to resolve the orbit at
+          n_cells (int, optional): number of grid cells to divide the calculation of occlusion to. Defaults to 20.
+
+      Returns:
+          Tuple[MeshModel, MeshModel]: (body1, body2)
+      """
+      grid = Grid.construct(binary.body1, binary.body2, n_cells)
+      return _evaluate_orbit(binary, time, grid)
+
+
+def evaluate_orbit_at_times(binary: Binary, times: ArrayLike, n_cells: int = 20) -> Tuple[MeshModel, MeshModel]:
+      """_summary_
+
+      Args:
+          binary (Binary): a binary object
+          times (ArrayLike): times to resolve the orbit at
+          n_cells (int, optional): number of grid cells to divide the calculation of occlusion to. Defaults to 20.
+
+      Returns:
+          Tuple[MeshModel, MeshModel]: (body1, body2)
+      """
+      grid = Grid.construct(binary.body1, binary.body2, n_cells)
+      return v_evaluate_orbit(binary, times, grid)