stellar-spice 0.2__py3-none-any.whl

spice/geometry/__init__.py

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

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))

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]])

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

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)

spice/models/mesh_generation.py

@@ -0,0 +1,214 @@
+from typing import Tuple
+from jax.typing import ArrayLike
+import jax.numpy as jnp
+import jax
+
+# https://sinestesia.co/blog/tutorials/python-icospheres/
+
+PHI: float = (1 + jnp.sqrt(5)) / 2
+
+def vertex(x: float, y: float, z: float) -> float:
+    """ Return vertex coordinates fixed to the unit sphere """
+    length = jnp.sqrt(jnp.power(x, 2) + jnp.power(y, 2) + jnp.power(z, 2))
+    coords = jnp.array([x, y, z])
+    return (coords)/length
+
+# Number of vertices and faces, solved in wolfram
+
+# 12 + 3*(faces-1) = 
+# faces: 4*(faces-1)
+
+# faces: f(n)=5*4^(n+1)
+# vertices: v(n) = 4*(5*4^n-2)
+
+# 0: f(0) = 4*5 = 20
+# 0: v(0) = 4*(5-2) = 12
+
+# 1: f(1) = 5*4^2 = 80
+# 1: v(1) = 4*(20-2) = 72
+
+def init_vertices_faces(subdiv: int) -> Tuple[ArrayLike, ArrayLike, ArrayLike]:
+    
+    init_verts = jnp.array([
+          vertex(-1,  PHI, 0),
+          vertex( 1,  PHI, 0),
+          vertex(-1, -PHI, 0),
+          vertex( 1, -PHI, 0),
+
+          vertex(0, -1, PHI),
+          vertex(0,  1, PHI),
+          vertex(0, -1, -PHI),
+          vertex(0,  1, -PHI),
+
+          vertex( PHI, 0, -1),
+          vertex( PHI, 0,  1),
+          vertex(-PHI, 0, -1),
+          vertex(-PHI, 0,  1),
+        ])
+
+    init_faces = jnp.array([
+             # 5 faces around point 0
+             [0, 11, 5],
+             [0, 5, 1],
+             [0, 1, 7],
+             [0, 7, 10],
+             [0, 10, 11],
+
+             # Adjacent faces
+             [1, 5, 9],
+             [5, 11, 4],
+             [11, 10, 2],
+             [10, 7, 6],
+             [7, 1, 8],
+
+             # 5 faces around 3
+             [3, 9, 4],
+             [3, 4, 2],
+             [3, 2, 6],
+             [3, 6, 8],
+             [3, 8, 9],
+
+             # Adjacent faces
+             [4, 9, 5],
+             [2, 4, 11],
+             [6, 2, 10],
+             [8, 6, 7],
+             [9, 8, 1],
+    ])
+    
+    verts_total_size = (4*(5*jnp.power(4, subdiv)-2)).astype(int)
+    verts = jnp.zeros((verts_total_size, 3))
+    verts = verts.at[:12].set(init_verts)
+    verts_mask = jnp.concatenate([jnp.ones(12,), jnp.zeros(verts_total_size-12,)]).astype(jnp.int32)
+    
+    faces_total_size = int(5*4**(subdiv+1))
+    faces = -1.*jnp.ones((faces_total_size, 3))
+    faces = faces.at[::int(faces_total_size/20)].set(init_faces)
+
+    return verts, verts_mask, faces
+
+@jax.jit
+def middle_point(point_1: int,
+                 point_2: int,
+                 verts: ArrayLike,
+                 verts_mask: ArrayLike,
+                 middle_point_cache: ArrayLike) -> Tuple[float,
+                                                         ArrayLike,
+                                                         ArrayLike, 
+                                                         ArrayLike]:
+    """ Find a middle point and project to the unit sphere """
+    # We check if we have already cut this edge first
+    # to avoid duplicated verts
+    smaller_index, greater_index = jax.lax.cond(point_1<point_2,
+                                                lambda: (point_1.astype(int), point_2.astype(int)),
+                                                lambda: (point_2.astype(int), point_1.astype(int)))
+    
+    def expand_verts(verts, verts_mask, point_1, point_2, middle_point_cache):
+        # If it's not in cache, then we can cut it
+        loc_verts = jnp.array([verts[point_1.astype(int)],
+                               verts[point_2.astype(int)]]).T
+        middle = jnp.sum(loc_verts, axis=1)/2
+
+        index = jnp.sum(verts_mask)
+        verts = verts.at[index].set(vertex(*middle))
+        verts_mask = verts_mask.at[index].set(1)
+        
+        smaller_index, greater_index = jax.lax.cond(point_1<point_2,
+                                                lambda: (point_1.astype(int), point_2.astype(int)),
+                                                lambda: (point_2.astype(int), point_1.astype(int)))
+
+        middle_point_cache = middle_point_cache.at[smaller_index, greater_index].set(index)
+
+        return index, verts, verts_mask, middle_point_cache
+    
+    maybe_index = middle_point_cache[smaller_index, greater_index]
+    return jax.lax.cond(maybe_index>0,
+                        lambda: (maybe_index, verts, verts_mask, middle_point_cache),
+                        lambda: expand_verts(verts, verts_mask, point_1, point_2, middle_point_cache))
+
+@jax.jit
+def subdivide_trie(carry: Tuple[ArrayLike, ArrayLike, ArrayLike], tri: ArrayLike) -> Tuple[ArrayLike,
+                                                                                           ArrayLike,
+                                                                                           ArrayLike,
+                                                                                           ArrayLike]:
+    verts, verts_mask, middle_point_cache = carry
+    
+    def _subdivide(verts, verts_mask, middle_point_cache, tri):
+        v1, verts, verts_mask, middle_point_cache = middle_point(tri[0],
+                                                                     tri[1],
+                                                                     verts,
+                                                                     verts_mask,
+                                                                     middle_point_cache)
+        v2, verts, verts_mask, middle_point_cache = middle_point(tri[1],
+                                                                     tri[2],
+                                                                     verts,
+                                                                     verts_mask,
+                                                                     middle_point_cache)
+        v3, verts, verts_mask, middle_point_cache = middle_point(tri[2],
+                                                                     tri[0],
+                                                                     verts,
+                                                                     verts_mask,
+                                                                     middle_point_cache)
+
+        return (verts, verts_mask, middle_point_cache), jnp.array([
+            [tri[0], v1, v3],
+            [tri[1], v2, v1],
+            [tri[2], v3, v2],
+            [v1, v2, v3]])
+    
+    return jax.lax.cond(jnp.all(tri==-1),
+                        lambda: ((verts, verts_mask, middle_point_cache), -1.*jnp.ones((4, 3))),
+                        lambda: _subdivide(verts, verts_mask, middle_point_cache, tri[0]))
+
+
+def fill_placeholders(x: jnp.array, chunk_size):
+    new_x = -1.*jnp.ones((chunk_size, 3))
+    return new_x.at[0].set(x)
+relax = jax.vmap(lambda x, n: fill_placeholders(x, n), in_axes=(0, None))
+
+
+@jax.jit
+def subdivide(faces: ArrayLike, verts: ArrayLike, verts_mask: ArrayLike) -> Tuple[ArrayLike,
+                                                                                  ArrayLike,
+                                                                                  ArrayLike]:
+    verts_total_size = verts.shape[0]
+    keys_cache = (-1*jnp.ones((verts_total_size, verts_total_size))).astype(int)
+    (verts, verts_mask, _), new_faces = jax.lax.scan(subdivide_trie,
+                                                     (verts, verts_mask, keys_cache),
+                                                     faces.reshape((-1, 4, 3)))
+    return new_faces.reshape((-1, 3)), verts, verts_mask
+
+
+relax_all = lambda chunks, n: jax.vmap(lambda x: relax(x[:4], int(n/4)), in_axes=(0,))(chunks.reshape((-1, n, 3))).reshape((-1, 3))
+
+@jax.jit
+def face_center(verts: ArrayLike, face: ArrayLike) -> Tuple[ArrayLike, float]:
+    a, b, c = verts[face[0]], verts[face[1]], verts[face[2]]
+    ab = b-a
+    ac = c-a
+    A = jnp.linalg.norm(jnp.cross(ab, ac))/2
+    return A, (a+b+c)/3
+
+def _icosphere(subdiv: int) -> Tuple[ArrayLike, ArrayLike, ArrayLike, ArrayLike]:
+    verts, verts_mask, faces = init_vertices_faces(subdiv)
+    
+    for s in range(subdiv):
+        faces, verts, verts_mask = subdivide(faces, verts, verts_mask)
+        current_s = int(faces.shape[0]/(5*4**(s+1)))
+        faces = relax_all(faces, current_s)
+    
+    areas, centers = jax.jit(jax.vmap(face_center, in_axes=(None, 0)))(verts, faces.astype(jnp.int32))
+    return verts, faces, areas, centers
+
+def icosphere(points: int) -> Tuple[ArrayLike, ArrayLike, ArrayLike, ArrayLike]:
+    """Create an icosphere with at least that n points.
+
+    Args:
+        points (int): Minimal number of vertices in the icosphere
+
+    Returns:
+        Tuple[ArrayLike, ArrayLike, ArrayLike, ArrayLike]: vertices (n, 3), faces (n, 3), triangle areas (n,), centers (n, 3)
+    """
+    subdivs = jnp.ceil(.5*jnp.log2(points/5)-1).astype(int)
+    verts, faces, areas, centers = _icosphere(subdivs)
+    return verts, faces, areas, centers