Python API

​

Components

Components are containers of data that is associated with an entity. See ECS Data Model for more context on entities and components.

To define a new component, add elodin.Component as metadata to a base class using typing.Annotated. The base class can be jax.Array or some other container of array data. This is an example of a component that annotates jax.Array:

import elodin as el

Wind = typing.Annotated[
    jax.Array,
    el.Component(
        "wind",
        el.ComponentType(el.PrimitiveType.F64, (3,)),
        metadata={"element_names": "x,y,z"},
    ),
]

​

class elodin.Component

A container of component metadata.

• __init__(name, type = None, asset = False, metadata = {}) -> elodin.Component

  Create a new component with:

  • Unique name (e.g. “world_pos, “inertia”).
  • Component type information (via elodin.ComponentType). This is optional if the base class already provides component type information as part of __metadata__, which is the case for elodin.Quaternion, elodin.Edge, and all spatial vector algebra classes.
  • Flag indicating whether the component is an asset (e.g. a mesh, texture, etc.).
  • Other metadata that is optional (e.g. description, units, labels, etc.).

  import elodin as el
  
  el.Component(
      "wind",
      el.ComponentType(el.PrimitiveType.F64, (3,)),
      metadata={"element_names": "x,y,z"},
  ),
  

  The above example defines a “wind” component that is a 3D vector of float64 values. The “element_names” entry is an example of optional metadata. It specifies the labels for each element of the vector that are displayed in the component inspector.

​

class elodin.ComponentType

ComponentType describes the shape and data type of a component. The shape is a tuple of integers that specifies the size of each dimension (e.g. () for scalars, (3,) for 3D vectors). The data type is an elodin.PrimitiveType.

• __init__(dtype, shape) -> elodin.ComponentType

  Create a component type from a data type and shape.

​

class elodin.Edge

An edge is a relationship between two entities. See elodin.GraphQuery for information on how to use edges in graph queries.

• __init__(left, right) -> elodin.Edge

  Create an edge between two entities given their unique ids.

​

Archetypes

An archetype is a combination of components with a unique name. To define a new archetype, create a subclass of elodin.Archetype with the desired components as fields. Here is an example of an archetype for a kalman filter:

import elodin as el
from dataclasses import dataclass

@dataclass
class KalmanFilter(el.Archetype):
    p: P
    att_est: AttEst
    ang_vel_est: AngVelEst
    bias_est: BiasEst

The archetype can then be used to attach components to entities:

world.insert(
    satellite,
    KalmanFilter(
        p=np.identity(6),
        att_est=el.Quaternion.identity(),
        ang_vel_est=np.zeros(3),
        bias_est=np.zeros(3),
    ),
)

​

Systems

Systems are the building blocks of simulation; they are functions that operate on a set of input components and produce a set of output components. Elodin provides decorators that allow for systems to be easily defined from functions.

​

@elodin.system

This is a lower-level API for defining a system. A function decorated with @elodin.system accepts special parameter types (such as elodin.Query and elodin.GraphQuery) that specify what data the system needs access to. It returns an elodin.Query containing one or more components. Some examples of @elodin.system are:

import elodin as el

@el.system
def gravity(
    graph: el.GraphQuery[GravityEdge],
    query: el.Query[el.WorldPos, el.Inertia],
) -> el.Query[el.Force]: ...

@el.system
def apply_wind(
    w: el.Query[Wind], q: el.Query[el.Force, el.WorldVel]
) -> el.Query[el.Force]: ...

​

@elodin.map

This is a higher-level API for defining a system that reduces the boilerplate of @elodin.system by unpacking the input and output queries into individual components, and wrapping the body of the function in a query.map(ret_type, ...) call. It is useful for systems with simple data flow patterns. Some examples of @elodin.map are:

import elodin as el

@el.map
def gravity(f: el.Force, inertia: el.Inertia) -> el.Force: ...

@el.map
def gyro_omega(vel: el.WorldVel) -> GyroOmega: ...

The following systems are equivalent as the @elodin.map definition effectively desugars to the @elodin.system one:

import elodin as el

@el.map
def gravity(f: el.Force, inertia: el.Inertia) -> el.Force:
    return f + el.SpatialForce.from_linear(inertia.mass() * jnp.array([0.0, -9.81, 0.0]))

@el.system
def gravity(query: el.Query[el.Force, el.Inertia]) -> el.Query[el.Force]:
    return query.map(
        el.Force,
        lambda f, inertia: f + el.SpatialForce.from_linear(inertia.mass() * jnp.array([0.0, -9.81, 0.0])),
    )

​

class elodin.Query

Query is the primary mechanism for accessing data in Elodin. It is a view into the world state that is filtered by the components specified in the query. Only entities that have been spawned with all of the query’s components will be selected for processing. For example, the query Query[WorldPos, Inertia] would only select entities that have both a WorldPos and an Inertia component (typically via the Body archetype).

• map(ret_type, map_fn) -> elodin.Query

  Apply a function map_fn to the query’s components and return a new query with the specified ret_type return type. map_fn should be a function that takes the query’s components as arguments and returns a single value of type ret_type.

  import elodin as el
  
  @el.system
  def gravity(query: el.Query[el.Force, el.Inertia]) -> el.Query[el.Force]:
      return query.map(
          el.Force,
          lambda f, inertia: f + el.SpatialForce.from_linear(inertia.mass() * jnp.array([0.0, -9.81, 0.0])),
      )
  

  In this example, ret_type is el.Force and map_fn is a lambda function with the signature (el.Force, el.Inertia) -> el.Force.

  To return multiple components as output, ret_type must be a tuple:

  import elodin as el
  
  @el.system
  def multi_out_sys(query: el.Query[A]) -> el.Query[C, D]:
      return query.map(tuple[C, D], lambda a: ...)
  

​

class elodin.GraphQuery

GraphQuery is a special type of query for operating on edges in an entity graph. Edges represent relationships between entities and are fundamental for modeling physics systems such as gravity.

A GraphQuery requires exactly one type argument, which must be an annotated elodin.Edge component. For example, GraphQuery[GravityEdge] is a valid graph query iff GravityEdge is a component with Edge as the base class:

GravityEdge = typing.Annotated[elodin.Edge, elodin.Component("gravity_edge")]

• edge_fold(left_query, right_query, ret_type, init_val, fold_fn) -> elodin.Query

  For each edge, query the left and right entity components using left_query and right_query, respectively. Then, apply the fold_fn function to those input components to compute the ret_type output component(s).

  The ret_type component(s) must belong to the left entity of the edge.

  A single left entity may have edges to multiple right entities, but it can only hold a single value for each ret_type component. So, the fold_fn computations for each entity’s edges must be accumulated into a single final value. To carry the intermediate results, fold_fn takes an “accumulator” value as the first argument. Its output is set as the accumulator value for the next iteration. init_val is the initial value of the accumulator.

  edge_fold makes no guarantees about the order in which edges are processed. For associative operators like +, the order the elements are combined in is not important, but for non-associative operators like -, the order will affect the final result.

  See the Three-Body Orbit Tutorial for a practical example of using edge_fold to compute gravitational forces between entities.

​

Primitives

​

class elodin.PrimitiveType

• elodin.PrimitiveType.F64 -> elodin.PrimitiveType

  A constant representing the 64-bit floating point data type.

• elodin.PrimitiveType.U64 -> elodin.PrimitiveType

  A constant representing the 64-bit unsigned integer data type.

​

class elodin.Quaternion

Unit quaternions are used to represent spatial orientations and rotations of bodies in 3D space.

• Quaternion.identity() -> elodin.Quaternion

  Create a unit quaternion with no rotation.

• Quaternion.from_axis_angle() -> elodin.Quaternion

  Create a quaternion from an axis and an angle.

• inverse() -> elodin.Quaternion

  Compute the inverse of the quaternion.

• normalize() -> elodin.Quaternion

  Normalize to a unit quaternion.

• __add__(other) -> elodin.Quaternion

  Add two quaternions.

  Adding quaternions does not yield the composite rotation unless they are infinitesimal rotations. Use multiplication instead.

• __mul__(other) -> elodin.Quaternion

  Multiply two quaternions.

• __matmul__(vector) -> jax.Array | elodin.SpatialTransform | elodin.SpatialMotion | elodin.SpatialForce

  Rotate vector by computing the matrix product. The vector can be a plain jax.Array or one of the following spatial objects: elodin.SpatialTransform, elodin.SpatialMotion, elodin.SpatialForce. The return type is the same as the input type.

​

Spatial Vector Algebra

Elodin uses Featherstone’s spatial vector algebra notation for rigid-body dynamics as it is a compact way of representing the state of a rigid body with six degrees of freedom. You can read a short into here or in Rigid Body Dynamics Algorithms (Featherstone - 2008).

​

class elodin.SpatialTransform

A spatial transform is a 7D vector that represents a rigid body transformation in 3D space.

• __init__(arr, angular, linear) -> elodin.SpatialTransform

  Create a spatial transform from either arr or angular and linear. If no arguments are provided, the spatial transform is initialized to the default values of the identity quaternion and the zero vector.

  • arr : jax.Array with shape (7)
  • angular : elodin.Quaternion, default is Quaternion.identity()
  • linear : jax.Array with shape (3), default is [0, 0, 0]

• linear() -> jax.Array

  Get the linear part of the spatial transform as a vector with shape (3,).

• angular() -> elodin.Quaternion

  Get the angular part of the spatial transform as a quaternion.

• __add__(other) -> elodin.SpatialTransform

  Add two spatial transforms.

​

class elodin.SpatialMotion

A spatial motion is a 6D vector that represents either the velocity or acceleration of a rigid body in 3D space.

• __init__(angular, linear) -> elodin.SpatialMotion

  Create a spatial motion from an angular and a linear vector. Both arguments are optional and default to zero vectors.

  • angular : jax.Array with shape (3), default is [0, 0, 0]
  • linear : jax.Array with shape (3), default is [0, 0, 0]

• linear() -> jax.Array

  Get the linear part of the spatial motion as a vector with shape (3,).

• angular() -> jax.Array

  Get the angular part of the spatial motion as a vector with shape (3,).

• __add__(other) -> elodin.SpatialMotion

  Add two spatial motions.

​

class elodin.SpatialForce

A spatial force is a 6D vector that represents the linear force and torque applied to a rigid body in 3D space.

• __init__(arr, torque, force) -> elodin.SpatialForce

  Create a spatial force from either arr or torque and force. If no arguments are provided, the spatial force is initialized to zero torque and force.

  • arr : jax.Array with shape (6)
  • torque : jax.Array with shape (3), default is [0, 0, 0]
  • force : jax.Array with shape (3), default is [0, 0, 0]

• force() -> jax.Array

  Get the linear force part of the spatial force as a vector with shape (3,).

• torque() -> jax.Array

  Get the torque part of the spatial force as a vector with shape (3,).

• __add__(other) -> elodin.SpatialForce

  Add two spatial forces.

​

class elodin.SpatialInertia

A spatial inertia is a 7D vector that represents the mass, moment of inertia, and momentum of a rigid body in 3D space. The moment of inertia is represented in its diagonalized form of $[I_1, I_2, I_3]$.

• __init__(mass, inertia) -> elodin.SpatialInertia

  Create a spatial tensor inertia from a scalar mass and an optional inertia tensor diagonal with shape (3,). If the inertia tensor is not provided, it is set to the same value as the mass along all axes.

• mass() -> jax.Array

  Get the scalar mass of the spatial inertia.

• inertia_diag() -> jax.Array

  Get the inertia tensor diagonal of the spatial inertia with shape (3,).