From 4640da759f38425f5794b9875d845443f479ea37 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20St=C3=B6lzle?= Date: Mon, 3 Aug 2026 20:50:51 -0400 Subject: [PATCH 1/5] Harmonize PCS and GVS parameters --- docs/api/actuation/index.md | 4 +- docs/api/systems/continuum-components.md | 199 ++++++ docs/api/systems/gvs/gvs.md | 68 +- docs/api/systems/gvs/index.md | 7 +- docs/api/systems/index.md | 7 + docs/api/systems/pcs/index.md | 5 + docs/api/systems/pcs/isupport.md | 33 +- docs/api/systems/pcs/pcs.md | 44 ++ docs/api/utilities/parameters.md | 108 +-- docs/development/contributing.md | 4 +- docs/development/extending.md | 21 +- docs/index.md | 22 +- docs/installation.md | 41 +- docs/user-guide/examples.md | 18 +- .../user-guide/parameters-and-optimization.md | 349 ++++++++++ docs/user-guide/quick-start.md | 60 +- .../setpoint_regulation_comparison.py | 25 +- ...ol_tendon_actuated_pcs_with_synergistic.py | 25 +- examples/simulation/gvs/simulate_gvs.py | 77 +-- .../gvs/simulate_tendon_actuated_gvs.py | 49 +- .../simulate_batched_tendon_actuated_pcs.py | 25 +- examples/simulation/pcs/simulate_isupport.py | 55 +- examples/simulation/pcs/simulate_pcs.py | 25 +- .../simulation/pcs/simulate_planar_pcs.py | 25 +- .../pcs/simulate_tendon_actuated_pcs.py | 29 +- .../simulate_tendon_actuated_planar_pcs.py | 23 +- mkdocs.yml | 2 + .../code/soromox/simulate_complex_gvs.py | 58 +- .../code/soromox/simulate_planar_pcs.py | 29 +- .../code/soromox/simulate_spatial_pcs.py | 29 +- .../soromox/simulate_tendon_driven_gvs.py | 48 +- .../code/identify_soft_tentacle_parameters.py | 89 +-- .../code/identify_soft_tentacle_residual.py | 45 +- ...nfiguration_space_comparison_simulation.py | 31 +- ...ompare_impedance_feedback_linearization.py | 27 +- .../operational_space_impedance_common.py | 31 +- ...ntrol_gain_optimization_with_collocated.py | 26 +- ...trol_gain_optimization_with_synergistic.py | 26 +- .../code/pcs_cf_cbf_clf_common.py | 34 +- .../code/parallel_soromox_env.py | 31 +- .../secVf_parallel_rl/code/render_rl_video.py | 26 +- src/soromox/rendering/open3d_renderer.py | 3 +- .../rendering/opencv_planar_renderer.py | 3 +- src/soromox/systems/__init__.py | 31 +- .../articulated/articulated_soft_robot.py | 3 +- src/soromox/systems/components/__init__.py | 31 + .../systems/components/cross_sections.py | 159 +++++ src/soromox/systems/components/joints.py | 335 ++++++++++ src/soromox/systems/components/links.py | 385 +++++++++++ src/soromox/systems/components/materials.py | 103 +++ src/soromox/systems/gvs/__init__.py | 10 +- src/soromox/systems/gvs/_assembly.py | 112 ++-- src/soromox/systems/gvs/_runtime.py | 52 +- src/soromox/systems/gvs/construction.py | 372 +++++++++-- src/soromox/systems/gvs/core.py | 628 +++++++++-------- src/soromox/systems/gvs/params.py | 251 +++---- src/soromox/systems/gvs/specs.py | 632 ++---------------- src/soromox/systems/gvs/structures.py | 47 +- src/soromox/systems/hsa/planar_hsa.py | 3 +- src/soromox/systems/params.py | 105 ++- src/soromox/systems/pcs/isupport.py | 163 ++--- src/soromox/systems/pcs/params.py | 147 ++-- src/soromox/systems/pcs/pcs.py | 603 ++++++++++++----- src/soromox/systems/pcs/planar_pcs.py | 566 ++++++++++++---- src/soromox/systems/pcs/structures.py | 26 +- src/soromox/systems/pendulum/pendulum.py | 3 +- src/soromox/systems/soft_robot.py | 8 - tests/actuation/test_threadlike.py | 39 +- tests/rendering/test_base_renderer.py | 2 +- .../rendering/test_isupport_viser_renderer.py | 31 +- .../rendering/test_open3d_material_frames.py | 2 +- tests/system_param_builders.py | 81 ++- tests/systems/test_gvs.py | 299 ++++----- tests/systems/test_pcs.py | 48 +- tests/systems/test_planar_pcs.py | 30 +- .../test_pressure_actuated_pcs_models.py | 117 +++- .../test_shared_continuum_components.py | 246 +++++++ tests/systems/test_soft_robot_defaults.py | 3 +- tests/systems/test_system_lengths.py | 28 +- tests/systems/test_typed_params_api.py | 215 +----- tools/benchmarks/_benchmark_common.py | 90 ++- 81 files changed, 5113 insertions(+), 2779 deletions(-) create mode 100644 docs/api/systems/continuum-components.md create mode 100644 docs/user-guide/parameters-and-optimization.md create mode 100644 src/soromox/systems/components/__init__.py create mode 100644 src/soromox/systems/components/cross_sections.py create mode 100644 src/soromox/systems/components/joints.py create mode 100644 src/soromox/systems/components/links.py create mode 100644 src/soromox/systems/components/materials.py create mode 100644 tests/systems/test_shared_continuum_components.py diff --git a/docs/api/actuation/index.md b/docs/api/actuation/index.md index 5da3cadf..5c938bcb 100644 --- a/docs/api/actuation/index.md +++ b/docs/api/actuation/index.md @@ -99,10 +99,10 @@ the physics modules to rendering types or styles. ## Parameter updates -Body parameters keep the existing API: +Body parameters use top-level or component-specific immutable updates: ```python -robot = robot.update_params(young_modulus=new_modulus) +robot = robot.update_link_params(stiffness=new_stiffness) robot = robot.with_params(new_body_params) ``` diff --git a/docs/api/systems/continuum-components.md b/docs/api/systems/continuum-components.md new file mode 100644 index 00000000..3385e75d --- /dev/null +++ b/docs/api/systems/continuum-components.md @@ -0,0 +1,199 @@ +# Continuum Robot Components + +PCS and GVS use the same public vocabulary for continuum links, cross-sections, +isotropic materials, and joints. This page describes the physical meaning and +ownership of those shared components. For construction, immutable replacement, +and gradient-based identification workflows, see +[Parameters and Optimization](../../user-guide/parameters-and-optimization.md). + +## Model composition + +Continuum-system parameters use shallow composition: + +```text +PCSParams / PlanarPCSParams +└── link: ContinuumLinkParams + └── cross_section: CrossSectionParams + +GVSParams +├── link: ContinuumLinkParams +│ └── cross_section: CrossSectionParams +└── joint: JointParams +``` + +The hierarchy expresses physical ownership rather than a different material +model for every system. A link owns its geometry, reference strain, mass +properties, stiffness, and damping. A GVS joint separately owns joint +stiffness and damping. System-level parameters own the base pose and gravity. + +## Cross-sections + +`CrossSectionParams.coefficients` is a two-dimensional array with one row per +link. The static system structure determines how each row is interpreted: + +- a scalar dimension defines a constant profile; +- `LinearProfile(base, tip)` defines a base-to-tip linear profile; +- rows are zero-padded when different GVS links require different numbers of + coefficients. + +The supported solid cross-sections are circular, rectangular, and elliptical. +`CrossSectionGeometry` identifies the family, while shared geometry utilities +calculate area and second moments of area. + +`LinkSpec` provides the ergonomic construction interface: + +```python +from soromox.systems import LinearProfile, LinkSpec + +circular = LinkSpec.circular( + length=0.2, + radius=0.012, + density=1000.0, + young_modulus=1.0e6, + shear_modulus=3.4e5, + material_damping_coefficient=1.0e4, + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], +) + +tapered_rectangle = LinkSpec.rectangular( + length=0.25, + height=LinearProfile(base=0.03, tip=0.02), + width=0.025, + density=1000.0, + young_modulus=1.0e6, + shear_modulus=3.4e5, + material_damping_coefficient=1.0e4, + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], +) +``` + +PCS currently accepts constant circular links. GVS accepts constant or linear +circular, rectangular, and elliptical cross-sections. + +## Continuum links + +`ContinuumLinkParams` stores the canonical numeric values used at runtime: + +| Field | Shape | Meaning | +|-------|-------|---------| +| `length` | `(N,)` | Link backbone lengths | +| `density` | `(N,)` | Volumetric mass densities | +| `reference_strain` | `(N, strain_dimension)` | Stress-free strain fields | +| `cross_section.coefficients` | `(N, max_coefficients)` | Packed geometry profiles | +| `stiffness` | `(N, generalized_dimension, generalized_dimension)` | Generalized link stiffness | +| `damping` | same as `stiffness` | Generalized link damping | + +Stiffness and damping are finite symmetric matrices. They are not required to +be diagonal or positive definite, which permits anisotropic, coupled, fitted, +and learned constitutive models. + +The generalized dimension depends on the system: + +| System | Canonical link matrix shape | +|--------|-----------------------------| +| Spatial PCS | `(N, 6, 6)` | +| PlanarPCS | `(N, 3, 3)` | +| GVS | `(N, max_dof, max_dof)` | + +GVS matrices are zero-padded beyond each link's active basis coordinates. + +## Isotropic material model + +Young's modulus, shear modulus, and material damping are useful construction +and identification variables, but they are not duplicated in the runtime robot +parameters. Instead, callers keep an `IsotropicMaterialParams` PyTree: + +```python +from soromox.systems import ( + IsotropicMaterialParams, + shear_modulus_from_poisson_ratio, +) + +young = 1.0e6 +material = IsotropicMaterialParams( + young_modulus=young, + shear_modulus=shear_modulus_from_poisson_ratio(young, 0.45), + material_damping_coefficient=1.0e4, +) +``` + +Each field may be scalar or contain one value per link. Scalar values are +broadcast when the material is applied. + +Geometry, link length, strain basis, rotational scaling, and quadrature are +projected into unit-response operators. Material matrices are then evaluated +as + +\[ +K_i = E_i K_{i,E} + G_i K_{i,G}, +\qquad +D_i = \eta_i D_{i,\eta}. +\] + +Equivalently, the unit operators represent the discretized integrals + +\[ +K_i = \int B_i(s)^\mathsf{T} C_i(s) B_i(s)\,\mathrm{d}s, +\qquad +D_i = \int B_i(s)^\mathsf{T} V_i(s) B_i(s)\,\mathrm{d}s. +\] + +`link_matrices_from_material` evaluates the mapping without modifying the +robot. `with_isotropic_material` returns a robot containing the resulting +canonical matrices. It does not store the material PyTree. + +Geometry updates refresh the unit-response operators but leave explicitly +stored matrices unchanged. Reapply `with_isotropic_material` when matrices +should follow updated material or geometry values. This prevents geometry +updates from silently overwriting an explicitly supplied constitutive model. + +## Joints + +GVS segments may start with a fixed, revolute, prismatic, helical, +cylindrical, planar, spherical, or free joint. `JointSpec` accepts stiffness +and damping in active joint coordinates. Omitted values create zero matrices. + +At runtime, `JointParams` stores both matrices with the same GVS padding as the +link matrices: + +```python +import jax.numpy as jnp + +from soromox.systems import JointSpec + +joint = JointSpec.revolute( + axis="z", + stiffness=jnp.array([[0.3]]), + damping=jnp.array([[0.02]]), +) +``` + +The GVS global matrices interleave joint and link contributions before +projection into active coordinates. PCS has no separate joint component. + +## Construction specs, params, and structure + +The three object categories have different lifetimes: + +| Category | Purpose | Typical examples | +|----------|---------|------------------| +| Specs | Convenient user input during construction | `LinkSpec`, `JointSpec`, `GVSSegment`, `StrainBasisSpec` | +| Params | Dynamic numeric JAX PyTrees | `ContinuumLinkParams`, `JointParams`, `PCSParams`, `GVSParams` | +| Structure | Static choices affecting layout or compilation | `PCSStructure`, `GVSStructure` | + +Specs may contain either isotropic material values or explicit generalized +matrices. Construction resolves them to one canonical runtime representation. +Changing numeric params with the same layout supports JAX transformations; +changing segment count, basis order, active strains, or padding requires +reconstruction with a new structure. + +## API reference + +::: soromox.systems.components + options: + show_root_heading: true + show_source: false + heading_level: 3 + group_by_category: true + docstring_section_style: table + members_order: source diff --git a/docs/api/systems/gvs/gvs.md b/docs/api/systems/gvs/gvs.md index b67c10de..03b85e0f 100644 --- a/docs/api/systems/gvs/gvs.md +++ b/docs/api/systems/gvs/gvs.md @@ -14,24 +14,42 @@ system-identification workflows. ## Quick Start +The shared link, cross-section, material, and joint model is described in +[Continuum Robot Components](../continuum-components.md). For full immutable +replacement and optimization workflows, see +[Parameters, Updates, and Optimization](../../../user-guide/parameters-and-optimization.md). + ```python import jax.numpy as jnp from soromox.systems import ( GVS, GVSSegment, JointSpec, + LinearProfile, LinkSpec, StrainBasisSpec, ) segment = GVSSegment( - link=LinkSpec.circular(E=1e6, nu=0.45, rho=1000.0, eta=1e4, L=0.3, r=0.03), - joint=JointSpec.fixed(), + link=LinkSpec.rectangular( + length=0.3, + height=LinearProfile(base=0.03, tip=0.02), + width=0.025, + density=1000.0, + young_modulus=1e6, + shear_modulus=3.45e5, + material_damping_coefficient=1e4, + reference_strain=[0, 0, 0, 1, 0, 0], + ), + joint=JointSpec.revolute( + axis="z", + stiffness=jnp.array([[0.3]]), + damping=jnp.array([[0.02]]), + ), basis=StrainBasisSpec( - type="monomial", - active=[1, 1, 1, 1, 0, 0], - orders=[1, 1, 1, 1, 0, 0], - xi_ref=[0, 0, 0, 1, 0, 0], + type="legendre", + strain_selector=("kappa_y", "sigma_x"), + basis_order=1, ), num_gauss_points=5, ) @@ -47,14 +65,14 @@ base_transform = robot.forward_kinematics(q, s=robot.segment_end_positions[-1]) `GVS.from_segments(...)` creates typed `GVSParams` and `GVSStructure` internally, so `robot.params` can be optimized or partially replaced later. Static structure contains no copied material constants, lengths, joint -stiffness, or reference strains. For workflows that need the split without +matrices, or reference strains. For workflows that need the split without constructing a robot, use `GVS.params_from_segments(...)`. ## Segment Specs -- `LinkSpec`: construction input for link geometry, material properties, and length. Its numeric values are copied into `GVSParams.link`; only the cross-section family remains static. -- `JointSpec`: construction input for joint type and optional axis, plane, pitch, and stiffness. Stiffness is copied into `GVSParams.joint_stiffness`. -- `StrainBasisSpec`: construction input for basis family, active strain components, basis orders, and reference strain. Reference strain is copied into `GVSParams.reference_strain`. +- `LinkSpec`: shared construction input for link geometry, reference strain, material properties or explicit generalized matrices, and length. +- `JointSpec`: shared construction input for joint type, kinematic choices, stiffness, and damping. Matrices are copied into `GVSParams.joint`. +- `StrainBasisSpec`: GVS-specific basis family, active strain components, and basis orders. - `GVSSegment`: combines one link, one preceding joint, one strain basis, and `num_gauss_points`. ## Basis And Joint Names @@ -75,7 +93,31 @@ After construction, GVS exposes canonical runtime arrays: - `dofs_per_segment`, `num_dofs`, `num_padded_dofs`, `active_dof_map` - `B_joint`, `B_Xs`, `B_Z1`, `B_Z2` - `xi_ref_joint`, `xi_ref_Xs`, `xi_ref_Z1`, `xi_ref_Z2` -- `mass_matrices`, `stiffness_matrices`, `damping_matrices`, `joint_stiffness` +- per-quadrature `mass_matrices` +- canonical `params.link.stiffness`, `params.link.damping` +- canonical `params.joint.stiffness`, `params.joint.damping` + +## Immutable updates + +```python +robot = robot.update_link_params( + stiffness=1.1 * robot.params.link.stiffness, + damping=0.9 * robot.params.link.damping, +) +robot = robot.update_joint_params( + damping=1.2 * robot.params.joint.damping, +) + +replacement = robot.params.replace( + link=robot.params.link.replace( + density=1.05 * robot.params.link.density, + ), + joint=robot.params.joint.replace( + stiffness=1.1 * robot.params.joint.stiffness, + ), +) +robot = robot.with_params(replacement) +``` ## When To Use GVS vs PCS @@ -102,13 +144,13 @@ After construction, GVS exposes canonical runtime arrays: show_source: false heading_level: 3 -::: soromox.systems.gvs.specs.LinkSpec +::: soromox.systems.components.links.LinkSpec options: show_root_heading: true show_source: false heading_level: 3 -::: soromox.systems.gvs.specs.JointSpec +::: soromox.systems.components.joints.JointSpec options: show_root_heading: true show_source: false diff --git a/docs/api/systems/gvs/index.md b/docs/api/systems/gvs/index.md index 6507bf8e..25bc1707 100644 --- a/docs/api/systems/gvs/index.md +++ b/docs/api/systems/gvs/index.md @@ -2,6 +2,11 @@ This section covers continuum soft robots modeled with the Geometric Variable Strain (GVS) approach. +GVS uses the shared continuum link, cross-section, isotropic material, and joint +model described in [Continuum Robot Components](../continuum-components.md). +See [Parameters, Updates, and Optimization](../../../user-guide/parameters-and-optimization.md) +for construction, replacement, and identification workflows. + ## Overview GVS systems generalize PCS by allowing arbitrary strain basis functions instead of piecewise constant assumptions. A model is assembled from `GVSSegment` entries, where each segment contains a `LinkSpec`, `JointSpec`, `StrainBasisSpec`, and `num_gauss_points`. @@ -13,7 +18,7 @@ GVS systems generalize PCS by allowing arbitrary strain basis functions instead | `GVSSegment` | Complete declaration of one GVS segment | | `LinkSpec` | Link geometry, material properties, and length | | `JointSpec` | Preceding joint type and optional joint parameters | -| `StrainBasisSpec` | Strain basis family, active components, orders, and reference strain | +| `StrainBasisSpec` | Strain basis family, active components, and orders | ## Basis Functions diff --git a/docs/api/systems/index.md b/docs/api/systems/index.md index 5e755b4d..9588e5d3 100644 --- a/docs/api/systems/index.md +++ b/docs/api/systems/index.md @@ -43,6 +43,13 @@ The `SoftRobot` class extends `DynamicalSystem` with interfaces specific to soft ## Common Components +PCS and GVS share continuum-link, cross-section, isotropic-material, and joint +components. Their physical meaning, ownership, matrix shapes, and construction +specifications are described in +[Continuum Robot Components](continuum-components.md). Practical immutable +replacement and gradient-based identification examples are in +[Parameters, Updates, and Optimization](../../user-guide/parameters-and-optimization.md). + ### SystemState The `SystemState` class is a container for the robot state used throughout SoRoMoX for simulation, control, and analysis. diff --git a/docs/api/systems/pcs/index.md b/docs/api/systems/pcs/index.md index 0bff21ff..07ec38ef 100644 --- a/docs/api/systems/pcs/index.md +++ b/docs/api/systems/pcs/index.md @@ -2,6 +2,11 @@ This section covers continuum soft robots modeled using the Piecewise Constant Strain (PCS) approach based on discrete Cosserat rod theory. +PCS uses the shared link, cross-section, and material model described in +[Continuum Robot Components](../continuum-components.md). See +[Parameters, Updates, and Optimization](../../../user-guide/parameters-and-optimization.md) +for construction, replacement, and identification workflows. + ## Overview PCS systems model continuum soft robots by dividing them into segments, each with constant strain. This approach provides: diff --git a/docs/api/systems/pcs/isupport.md b/docs/api/systems/pcs/isupport.md index 70f29a99..d70709d4 100644 --- a/docs/api/systems/pcs/isupport.md +++ b/docs/api/systems/pcs/isupport.md @@ -77,19 +77,30 @@ state callback that returns `(q, pressures)`. ```python import jax.numpy as jnp -from soromox.systems import ISupport, ISupportParams, ISupportStructure +from soromox.systems import ISupport, ISupportParams, ISupportStructure, LinkSpec, PCS -params = ISupportParams( +reference = [0.0, 0.0, 0.0, 1.0, 0.0, 0.0] +links = [ # Physical order: rigid base, pneumatic section, rigid tip. - length=jnp.array([0.01, 0.18, 0.01]), - radius=jnp.array([0.03, 35.6e-3, 0.025]), - density=jnp.array([1210.0, 1104.0, 1210.0]), - young_modulus=jnp.array([2.0e9, 1.6464e6, 2.0e9]), - shear_modulus=jnp.array([0.8e9, 0.5488e6, 0.8e9]), - material_damping_coefficient=1.96e3, - reference_strain=jnp.tile( - jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), 3 - ), + LinkSpec.circular( + length=length, + radius=radius, + density=density, + young_modulus=young, + shear_modulus=shear, + material_damping_coefficient=1.96e3, + reference_strain=reference, + ) + for length, radius, density, young, shear in zip( + [0.01, 0.18, 0.01], + [0.03, 35.6e-3, 0.025], + [1210.0, 1104.0, 1210.0], + [2.0e9, 1.6464e6, 2.0e9], + [0.8e9, 0.5488e6, 0.8e9], + ) +] +params = ISupportParams( + link=PCS.params_from_links(links).link, chamber_inner_radius=jnp.array([6.39e-3]), chamber_outer_radius=jnp.array([7.79e-3]), chamber_distance=jnp.array([20e-3]), diff --git a/docs/api/systems/pcs/pcs.md b/docs/api/systems/pcs/pcs.md index 5baa3047..834427da 100644 --- a/docs/api/systems/pcs/pcs.md +++ b/docs/api/systems/pcs/pcs.md @@ -6,6 +6,50 @@ The general Piecewise Constant Strain (PCS) implementation provides the core mod This module contains the fundamental PCS implementation that serves as the foundation for more specialized PCS variants. It provides the core mathematical framework for modeling continuum robots using piecewise constant strain assumptions, following the discrete Cosserat approach for multisection soft manipulator dynamics. +## Construction and updates + +The shared component model is described in +[Continuum Robot Components](../continuum-components.md). For full immutable +replacement and optimization workflows, see +[Parameters, Updates, and Optimization](../../../user-guide/parameters-and-optimization.md). + +```python +import jax.numpy as jnp + +from soromox.systems import LinkSpec, PCS + +pcs = PCS.from_links([ + LinkSpec.circular( + length=0.2, + radius=0.012, + density=1000.0, + young_modulus=1.0e6, + shear_modulus=3.4e5, + material_damping_coefficient=1.0e4, + reference_strain=[0, 0, 0, 1, 0, 0], + ) +]) + +# Explicit generalized matrices bypass isotropic material construction. +explicit = PCS.from_links([ + LinkSpec.circular( + length=0.15, + radius=0.01, + density=1000.0, + stiffness=jnp.diag(jnp.array([0.2, 0.8, 0.8, 100.0, 30.0, 30.0])), + damping=jnp.diag(jnp.array([0.01, 0.02, 0.02, 0.5, 0.1, 0.1])), + reference_strain=[0, 0, 0, 1, 0, 0], + ) +]) + +pcs = pcs.update_link_params( + stiffness=1.1 * pcs.params.link.stiffness, + damping=0.9 * pcs.params.link.damping, +) +replacement = pcs.params.link.replace(density=1.05 * pcs.params.link.density) +pcs = pcs.with_params(pcs.params.replace(link=replacement)) +``` + ## API Reference ::: soromox.systems.pcs.pcs diff --git a/docs/api/utilities/parameters.md b/docs/api/utilities/parameters.md index 5743292c..3f718158 100644 --- a/docs/api/utilities/parameters.md +++ b/docs/api/utilities/parameters.md @@ -1,13 +1,18 @@ # Parameters +Shared link, joint, cross-section, and isotropic-material mechanics are +described in [Continuum Robot Components](../systems/continuum-components.md). +Construction, immutable replacement, and gradient-based workflows are covered +in [Parameters, Updates, and Optimization](../../user-guide/parameters-and-optimization.md). + Actuator and passive-element parameters follow the same immutable replacement style through indexed robot delegates. See [Actuation parameter updates](../actuation/index.md#parameter-updates). -System parameters are represented as typed Equinox PyTrees. Shared base classes -and cross-system tendon params live in `soromox.systems.params`; concrete params -and structures live next to their system family, for example -`soromox.systems.gvs.params` and `soromox.systems.gvs.structures`. +System parameters are represented as typed Equinox PyTrees. Shared continuum +link, joint, cross-section, material params, and construction specs live in +`soromox.systems.components`. System params and static structures remain next +to their system family. ## Overview @@ -19,10 +24,10 @@ Each system separates dynamic numeric values from static model structure: strain masks, GVS joint/basis/cross-section choices, symbolic expression paths, and padding sizes. Changing structure means constructing a new system and may recompile jitted methods. -- **Spec objects** are ergonomic construction inputs for model families that - need richer setup. For GVS, `GVSSegment`, `LinkSpec`, `JointSpec`, and - `StrainBasisSpec` may contain both static choices and numeric values; factory - methods split them into params and structure objects. +- **Spec objects** are ergonomic construction inputs. Shared `LinkSpec` and + `JointSpec` objects live in `soromox.systems.components`; `GVSSegment` and + `StrainBasisSpec` remain GVS-specific. Factory methods split specs into + runtime params and static structures. The top-level `soromox.systems` package re-exports the public params, structures, and specs for convenient imports. Internally, concrete containers @@ -30,8 +35,9 @@ are family-local: | System family | Dynamic params | Static structures | Construction specs | |---------------|----------------|-------------------|--------------------| -| PCS | `soromox.systems.pcs.params` | `soromox.systems.pcs.structures` | - | -| GVS | `soromox.systems.gvs.params` | `soromox.systems.gvs.structures` | `soromox.systems.gvs.specs` | +| Shared components | `soromox.systems.components` | - | `LinkSpec`, `JointSpec` | +| PCS | `soromox.systems.pcs.params` | `soromox.systems.pcs.structures` | shared `LinkSpec` | +| GVS | `soromox.systems.gvs.params` | `soromox.systems.gvs.structures` | shared specs plus `soromox.systems.gvs.specs` | | HSA | `soromox.systems.hsa.params` | `soromox.systems.hsa.structures` | - | | Pendulum | `soromox.systems.pendulum.params` | - | - | | Articulated | `soromox.systems.articulated.params` | - | - | @@ -39,8 +45,8 @@ are family-local: The public construction pattern is: ```python -robot = PCS(params=PCSParams(...), structure=PCSStructure(...)) -robot = robot.update_params(length=new_length) +robot = PCS.from_links(links, structure=PCSStructure(...)) +robot = robot.update_link_params(length=new_length) robot = robot.with_params(new_params) ``` @@ -49,9 +55,10 @@ Changing the number of segments, tendons, active strains, GVS basis layout, or quadrature layout is a structural change and requires reconstruction. For GVS specifically, `GVS.from_segments(...)` is the recommended constructor. -It accepts user-facing segment specs, stores numeric values only in `GVSParams`, -and stores stripped static choices in `GVSStructure`. This avoids stale -duplicates when updating values such as Young's modulus or link length. +It accepts user-facing segment specs, stores canonical link and joint matrices +in `GVSParams`, and stores stripped static choices in `GVSStructure`. +`IsotropicMaterialParams` remains a separate caller-owned optimization PyTree, +so runtime and material representations are not duplicated. ## Naming @@ -63,17 +70,20 @@ or `(num_links,)`. | Field | Meaning | |-------|---------| | `length` | Per-segment or per-link length | -| `radius` | Per-segment circular cross-section radius | +| `cross_section.coefficients` | Batched cross-section profile coefficients | | `density` | Per-segment material density | -| `young_modulus` | Per-segment Young's modulus | -| `shear_modulus` | Per-segment shear modulus | -| `material_damping_coefficient` | PCS material damping coefficient | -| `damping_matrix` | Custom generalized damping matrix | +| `stiffness` | Canonical per-link or per-joint generalized stiffness matrices | +| `damping` | Canonical per-link or per-joint generalized damping matrices | | `gravity` | Gravity vector | | `base_pose` | Base configuration as scalar-first quaternion pose | | `reference_strain` | Reference strain vector | | `joint_rest_configuration` | Joint coordinates where elastic joint force is zero | +Young's modulus, shear modulus, and material damping use the separate +`IsotropicMaterialParams` fields `young_modulus`, `shear_modulus`, and +`material_damping_coefficient`. See the linked components page for their +mapping to canonical matrices and the user guide for optimization examples. + ## World Frame, Mounting, and Gravity Defaults Soft-robot parameter objects use an upright mounting and Earth gravity when @@ -95,31 +105,24 @@ The exact default values are: - Planar poses use `[theta, x, y]`. Spatial poses use scalar-first Hamilton quaternions in `[qw, qx, qy, qz, x, y, z]` order. -Omitting both fields selects the defaults: +Omitting both fields selects the defaults. For example, the factory below +constructs an upright PlanarPCS under Earth gravity: ```python -params = PlanarPCSParams( - length=length, - radius=radius, - density=density, - young_modulus=young_modulus, - shear_modulus=shear_modulus, - damping_matrix=damping_matrix, - reference_strain=reference_strain, -) +robot = PlanarPCS.from_links(planar_links) ``` Use the inherited mounting constructors to make another common mounting -explicit. `base_position` translates the mounting without changing its -orientation: +explicit when constructing a param tree directly. `base_position` translates +the mounting without changing its orientation: ```python -horizontal = PlanarPCSParams.horizontal(**planar_params) +horizontal = PlanarPCSParams.horizontal(link=planar_link_params) upright = PlanarPCSParams.upright( - **planar_params, base_position=jnp.array([0.2, 0.1]) + link=planar_link_params, base_position=jnp.array([0.2, 0.1]) ) hanging = PCSParams.hanging( - **spatial_params, base_position=jnp.array([0.0, 0.0, 0.5]) + link=spatial_link_params, base_position=jnp.array([0.0, 0.0, 0.5]) ) ``` @@ -127,9 +130,9 @@ Pass an explicit vector for custom or zero gravity. Pass an explicit `base_pose` through the ordinary constructor for arbitrary orientations: ```python -zero_gravity = PCSParams(..., gravity=jnp.zeros(3)) +zero_gravity = PCSParams(link=spatial_link_params, gravity=jnp.zeros(3)) custom = PlanarPCSParams( - ..., + link=planar_link_params, gravity=jnp.array([1.0, -9.7]), base_pose=jnp.array([0.3, 0.2, 0.1]), ) @@ -193,20 +196,23 @@ choices that affect compilation. ```python import jax.numpy as jnp -from soromox.systems import PCS, PCSParams, PCSStructure - -params = PCSParams.upright( - length=jnp.array([0.1, 0.1]), - radius=jnp.array([0.01, 0.01]), - density=jnp.array([1000.0, 1000.0]), - young_modulus=jnp.array([1e6, 1e6]), - shear_modulus=jnp.array([1e5, 1e5]), - damping_matrix=jnp.eye(12), - reference_strain=jnp.tile(jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), 2), -) - -robot = PCS(params=params, structure=PCSStructure(num_gauss_points=5)) -updated_robot = robot.update_params(length=jnp.array([0.12, 0.1])) +from soromox.systems import LinkSpec, PCS, PCSStructure + +links = [ + LinkSpec.circular( + length=0.1, + radius=0.01, + density=1000.0, + young_modulus=1e6, + shear_modulus=1e5, + material_damping_coefficient=1e4, + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + ) + for _ in range(2) +] + +robot = PCS.from_links(links, structure=PCSStructure(num_gauss_points=5)) +updated_robot = robot.update_link_params(length=jnp.array([0.12, 0.1])) ``` ## API Reference diff --git a/docs/development/contributing.md b/docs/development/contributing.md index c981e0ff..2315ff2b 100644 --- a/docs/development/contributing.md +++ b/docs/development/contributing.md @@ -174,7 +174,9 @@ def forward_kinematics(params: PCSParams, q: Array) -> Array: ValueError: If configuration vector has wrong dimensions Example: - >>> params = params.replace(length=jnp.array([0.1])) + >>> params = params.replace( + ... link=params.link.replace(length=jnp.array([0.1])) + ... ) >>> q = jnp.array([0.0, 0.0, -1.0]) >>> pos = forward_kinematics(params, q) """ diff --git a/docs/development/extending.md b/docs/development/extending.md index dd4149e9..b9d496bc 100644 --- a/docs/development/extending.md +++ b/docs/development/extending.md @@ -1063,20 +1063,19 @@ def error_based_feedback_term( #### 4. Usage Example ```python -from soromox.systems import PlanarPCS, PlanarPCSParams +from soromox.systems import LinkSpec, PlanarPCS from soromox.control import ReferenceTrajectory # Create robot -params = PlanarPCSParams( - length=jnp.array([0.1, 0.1, 0.1]), - radius=jnp.array([0.01, 0.01, 0.01]), - density=jnp.array([1000.0, 1000.0, 1000.0]), - young_modulus=jnp.array([1e6, 1e6, 1e6]), - shear_modulus=jnp.array([1e5, 1e5, 1e5]), - material_damping_coefficient=jnp.array([318.0, 318.0, 318.0]), - reference_strain=jnp.tile(jnp.array([0.0, 1.0, 0.0]), 3), -) -robot = PlanarPCS(params=params) +robot = PlanarPCS.from_links([ + LinkSpec.circular( + length=0.1, radius=0.01, density=1000.0, + young_modulus=1e6, shear_modulus=1e5, + material_damping_coefficient=318.0, + reference_strain=[0.0, 1.0, 0.0], + ) + for _ in range(3) +]) # Create reference trajectory (defines desired motion) ref_traj = ReferenceTrajectory(...) # Your trajectory diff --git a/docs/index.md b/docs/index.md index b3760fd9..043695f9 100644 --- a/docs/index.md +++ b/docs/index.md @@ -111,20 +111,20 @@ Get up and running in minutes: ```python import jax.numpy as jnp - from soromox.systems import PlanarPCS, PlanarPCSParams, SystemState + from soromox.systems import LinkSpec, PlanarPCS, SystemState # Create a planar PCS soft robot num_segments = 3 - params = PlanarPCSParams( - length=0.1 * jnp.ones((num_segments,)), - radius=0.02 * jnp.ones((num_segments,)), - density=1070.0 * jnp.ones((num_segments,)), - reference_strain=jnp.tile(jnp.array([0.0, 1.0, 0.0]), num_segments), - young_modulus=2e3 * jnp.ones((num_segments,)), - shear_modulus=1e3 * jnp.ones((num_segments,)), - material_damping_coefficient=318.0, - ) - robot = PlanarPCS(params=params) + links = [ + LinkSpec.circular( + length=0.1, radius=0.02, density=1070.0, + reference_strain=[0.0, 1.0, 0.0], + young_modulus=2e3, shear_modulus=1e3, + material_damping_coefficient=318.0, + ) + for _ in range(num_segments) + ] + robot = PlanarPCS.from_links(links) # Initialize state q0 = jnp.zeros(robot.n_q) diff --git a/docs/installation.md b/docs/installation.md index 9a06718f..df1b4af2 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -170,19 +170,15 @@ Test your installation with this quick verification script: ```python import jax.numpy as jnp - from soromox.systems import PlanarPCS, PlanarPCSParams + from soromox.systems import LinkSpec, PlanarPCS # Create a simple 1-segment PCS robot - params = PlanarPCSParams( - length=jnp.array([0.1]), - radius=jnp.array([0.01]), - density=jnp.array([1000.0]), - young_modulus=jnp.array([1e6]), - shear_modulus=jnp.array([1e5]), - material_damping_coefficient=jnp.array([318.0]), - reference_strain=jnp.array([0.0, 1.0, 0.0]), - ) - robot = PlanarPCS(params=params) + robot = PlanarPCS.from_links([LinkSpec.circular( + length=0.1, radius=0.01, density=1000.0, + young_modulus=1e6, shear_modulus=1e5, + material_damping_coefficient=318.0, + reference_strain=[0.0, 1.0, 0.0], + )]) robot.forward_kinematics(jnp.zeros(robot.num_dofs), s=0.1) print("🎉 SoRoMoX installation successful!") @@ -193,18 +189,17 @@ Test your installation with this quick verification script: ```python import jax import jax.numpy as jnp - from soromox.systems import PlanarPCS, PlanarPCSParams - - params = PlanarPCSParams( - length=jnp.array([0.1, 0.1]), - radius=jnp.array([0.01, 0.01]), - density=jnp.array([1000.0, 1000.0]), - young_modulus=jnp.array([1e6, 1e6]), - shear_modulus=jnp.array([1e5, 1e5]), - material_damping_coefficient=jnp.array([318.0, 318.0]), - reference_strain=jnp.tile(jnp.array([0.0, 1.0, 0.0]), 2), - ) - robot = PlanarPCS(params=params) + from soromox.systems import LinkSpec, PlanarPCS + + robot = PlanarPCS.from_links([ + LinkSpec.circular( + length=0.1, radius=0.01, density=1000.0, + young_modulus=1e6, shear_modulus=1e5, + material_damping_coefficient=318.0, + reference_strain=[0.0, 1.0, 0.0], + ) + for _ in range(2) + ]) # Test JAX compilation and differentiation. Same-shape params updates keep # the PyTree layout fixed and avoid recompilation. diff --git a/docs/user-guide/examples.md b/docs/user-guide/examples.md index 4dc08daa..b002fb57 100644 --- a/docs/user-guide/examples.md +++ b/docs/user-guide/examples.md @@ -1,5 +1,8 @@ # Examples +See [Parameters, Updates, and Optimization](parameters-and-optimization.md) +for complete construction, immutable update, JAX, and Optax examples. + This page provides a comprehensive overview of the example scripts included with SoRoMoX. Examples are organized by application area to help you find the right starting point for your use case. !!! tip "Getting Started" @@ -273,14 +276,21 @@ Most examples allow easy parameter modification. Understanding parameter structu **For PCS systems:** ```python -# Modify physical parameters -params = params.replace( +# Modify canonical runtime link parameters +section = robot.params.link.cross_section.replace( + coefficients=0.9 * robot.params.link.cross_section.coefficients +) +robot = robot.update_link_params( length=0.2 * jnp.ones((num_segments,)), - radius=0.01 * jnp.ones((num_segments,)), + cross_section=section, +) + +# Material variables are caller-owned and applied explicitly. +material = material.replace( young_modulus=1e6 * jnp.ones((num_segments,)), shear_modulus=5e5 * jnp.ones((num_segments,)), ) -robot = robot.with_params(params) +robot = robot.with_isotropic_material(material) ``` **For pendulum systems:** diff --git a/docs/user-guide/parameters-and-optimization.md b/docs/user-guide/parameters-and-optimization.md new file mode 100644 index 00000000..b2588f4c --- /dev/null +++ b/docs/user-guide/parameters-and-optimization.md @@ -0,0 +1,349 @@ +# Parameters, Updates, and Optimization + +SoRoMoX separates convenient construction specifications, dynamic numeric +parameters, and static model structure. This guide shows how to construct PCS +and GVS systems, update their parameter PyTrees immutably, and use the same +interfaces in gradient-based optimization. + +For the mechanical meaning of links, joints, cross-sections, stiffness, +damping, and isotropic materials, see +[Continuum Robot Components](../api/systems/continuum-components.md). + +## Choosing the right parameter layer + +| Task | Interface | +|------|-----------| +| Construct a robot from physical descriptions | `LinkSpec`, `JointSpec`, `GVSSegment` | +| Replace runtime numeric values | `update_params`, `update_link_params`, `update_joint_params` | +| Replace several nested components together | nested `.replace(...)` followed by `with_params(...)` | +| Identify isotropic material properties | caller-owned `IsotropicMaterialParams` and `with_isotropic_material(...)` | +| Fit anisotropic or coupled mechanics | optimize `params.link.stiffness` and `params.link.damping` directly | +| Change basis order, segment count, or padding | reconstruct with a new structure | + +All update methods return new objects. They do not mutate the original robot or +parameter PyTree. + +## Constructing PCS and GVS + +The shared `LinkSpec` accepts either isotropic material properties or explicit +generalized matrices. This PCS example derives shear modulus from Poisson's +ratio: + +```python +import jax.numpy as jnp + +from soromox.systems import ( + PCS, + LinkSpec, + shear_modulus_from_poisson_ratio, +) + +young = 1.0e6 +shear = shear_modulus_from_poisson_ratio( + young, + poisson_ratio=0.45, +) +reference = [0.0, 0.0, 0.0, 1.0, 0.0, 0.0] + +pcs = PCS.from_links([ + LinkSpec.circular( + length=0.20, + radius=0.012, + density=1000.0, + young_modulus=young, + shear_modulus=shear, + material_damping_coefficient=1.0e4, + reference_strain=reference, + ) +]) +``` + +GVS uses the same link description and adds joint, strain-basis, and +quadrature specifications: + +```python +from soromox.systems import ( + GVS, + GVSSegment, + JointSpec, + StrainBasisSpec, +) + +gvs = GVS.from_segments([ + GVSSegment( + link=LinkSpec.circular( + length=0.20, + radius=0.012, + density=1000.0, + young_modulus=young, + shear_modulus=shear, + material_damping_coefficient=1.0e4, + reference_strain=reference, + ), + joint=JointSpec.revolute( + axis="z", + stiffness=jnp.array([[0.3]]), + damping=jnp.array([[0.02]]), + ), + basis=StrainBasisSpec( + type="legendre", + strain_selector=("kappa_y", "sigma_x"), + basis_order=1, + ), + num_gauss_points=7, + ) +]) +``` + +An explicit `stiffness=` or `damping=` on `LinkSpec` bypasses the corresponding +material construction: + +```python +explicit_pcs = PCS.from_links([ + LinkSpec.circular( + length=0.15, + radius=0.01, + density=1000.0, + stiffness=jnp.diag( + jnp.array([0.2, 0.8, 0.8, 100.0, 30.0, 30.0]) + ), + damping=jnp.diag( + jnp.array([0.01, 0.02, 0.02, 0.5, 0.1, 0.1]) + ), + reference_strain=reference, + ) +]) +``` + +## Immutable runtime updates + +Use the narrowest update method that owns a parameter. PCS, PlanarPCS, and GVS +share `update_link_params`: + +```python +stiffness = pcs.params.link.stiffness +stiffness = stiffness.at[0].set(1.1 * stiffness[0]) + +pcs = pcs.update_link_params( + stiffness=stiffness, + damping=0.9 * pcs.params.link.damping, +) +``` + +GVS additionally exposes joint updates: + +```python +joint_damping = gvs.params.joint.damping +joint_damping = joint_damping.at[0, 0, 0].set(0.03) +gvs = gvs.update_joint_params(damping=joint_damping) +``` + +Use nested `.replace(...)` and `with_params(...)` when replacing multiple +components atomically: + +```python +replacement = gvs.params.replace( + link=gvs.params.link.replace( + density=1.05 * gvs.params.link.density, + stiffness=1.1 * gvs.params.link.stiffness, + ), + joint=gvs.params.joint.replace( + stiffness=1.2 * gvs.params.joint.stiffness, + damping=1.1 * gvs.params.joint.damping, + ), +) +gvs = gvs.with_params(replacement) +``` + +Top-level values use `update_params`: + +```python +pcs = pcs.update_params(gravity=jnp.zeros(3)) +``` + +Updates may change numeric values but not static layout. Changing the number of +links, GVS basis orders, active strain selectors, joint types, matrix padding, +or quadrature padding requires reconstruction. + +## Isotropic material updates + +Generalized stiffness and damping matrices are the canonical runtime values. +Young's modulus, shear modulus, and material damping remain convenient +construction and optimization variables in a separate caller-owned PyTree: + +```python +from soromox.systems import IsotropicMaterialParams + +material = IsotropicMaterialParams( + young_modulus=jnp.array([1.0e6]), + shear_modulus=jnp.array([3.4e5]), + material_damping_coefficient=jnp.array([1.0e4]), +) + +material = material.replace( + young_modulus=1.1 * material.young_modulus, + shear_modulus=0.95 * material.shear_modulus, + material_damping_coefficient=( + 1.2 * material.material_damping_coefficient + ), +) +updated_pcs = pcs.with_isotropic_material(material) +``` + +The same call works for PlanarPCS and GVS. Scalar fields are broadcast; arrays +must contain one value per link. The robot stores only the generated canonical +matrices, not `material`. + +Inspect the mapping without updating the robot: + +```python +stiffness, damping = pcs.link_matrices_from_material(material) +updated = pcs.with_isotropic_material(material) + +assert jnp.allclose(updated.params.link.stiffness, stiffness) +assert jnp.allclose(updated.params.link.damping, damping) +``` + +## Differentiating parameter updates + +Typed params and robot updates are JAX PyTrees, so a loss can construct a +candidate robot without mutation. This example differentiates all three +isotropic material fields: + +```python +import jax + +robot = pcs.with_isotropic_material(material) +target_stiffness = 1.05 * robot.stiffness_matrix() +target_damping = 0.9 * robot.damping_matrix(jnp.zeros(robot.num_dofs)) + +def material_loss(candidate_material): + candidate = robot.with_isotropic_material(candidate_material) + stiffness_error = candidate.stiffness_matrix() - target_stiffness + damping_error = ( + candidate.damping_matrix(jnp.zeros(candidate.num_dofs)) + - target_damping + ) + return ( + jnp.mean(stiffness_error**2) + + jnp.mean(damping_error**2) + ) + +value, gradient = jax.value_and_grad(material_loss)(material) +print(value) +print(gradient.young_modulus) +print(gradient.shear_modulus) +print(gradient.material_damping_coefficient) +``` + +Cached unit-response operators reduce every material evaluation to batched +scalar-matrix combinations: + +```text +K = E K_E + G K_G +D = eta D_eta +``` + +The same immutable pattern applies to other numeric fields. For example, +canonical stiffness can be differentiated directly: + +```python +def matrix_loss(link_stiffness): + candidate = robot.update_link_params(stiffness=link_stiffness) + return jnp.mean( + (candidate.stiffness_matrix() - target_stiffness) ** 2 + ) + +matrix_value, matrix_gradient = jax.value_and_grad(matrix_loss)( + robot.params.link.stiffness +) +``` + +## Positive parameterization + +Optimize material properties in log space when they must remain nonnegative. +This constrains the material scalars without imposing positive-definiteness on +the canonical generalized matrices: + +```python +log_material = jax.tree.map(jnp.log, material) + +def decode(log_values): + return IsotropicMaterialParams( + young_modulus=jnp.exp(log_values.young_modulus), + shear_modulus=jnp.exp(log_values.shear_modulus), + material_damping_coefficient=jnp.exp( + log_values.material_damping_coefficient + ), + ) +``` + +## Optax optimization loop + +The following loop uses the project's examples dependency and keeps the +optimized variable as the log-material PyTree: + +```python +import optax + +optimizer = optax.adam(1.0e-2) +optimizer_state = optimizer.init(log_material) + +def log_loss(log_values): + candidate = robot.with_isotropic_material(decode(log_values)) + return jnp.mean( + (candidate.stiffness_matrix() - target_stiffness) ** 2 + ) + +value_and_grad = jax.jit(jax.value_and_grad(log_loss)) +for _ in range(100): + value, gradients = value_and_grad(log_material) + updates, optimizer_state = optimizer.update( + gradients, + optimizer_state, + log_material, + ) + log_material = optax.apply_updates(log_material, updates) + +final_material = decode(log_material) +optimized_robot = robot.with_isotropic_material(final_material) +``` + +`robot` may be a PCS, PlanarPCS, or GVS instance. + +## Geometry and material co-optimization + +Length or cross-section updates refresh unit-response operators but do not +silently overwrite canonical matrices. Apply geometry first and material +second: + +```python +new_section = robot.params.link.cross_section.replace( + coefficients=( + 1.05 * robot.params.link.cross_section.coefficients + ) +) +geometry_robot = robot.update_link_params(cross_section=new_section) + +assert jnp.allclose( + geometry_robot.params.link.stiffness, + robot.params.link.stiffness, +) +cooptimized_robot = geometry_robot.with_isotropic_material(material) +``` + +In a geometry/material optimization loop, construct `geometry_robot` from the +candidate geometry first, then apply the candidate material. This evaluates the +material mapping with the refreshed geometry operators while preserving the +explicit update order. + +## Material variables or direct matrices? + +Prefer isotropic material optimization when Young's modulus, shear modulus, and +material damping have physical meaning in the identification problem. Optimize +canonical `params.link.stiffness` and `params.link.damping` directly for +anisotropic, coupled, or learned constitutive models. + +PCS matrices have shape `(N, 6, 6)`, PlanarPCS matrices `(N, 3, 3)`, and GVS +matrices `(N, max_dof, max_dof)` with zero padding beyond each link's active +coordinates. The update and optimization pattern is otherwise identical. diff --git a/docs/user-guide/quick-start.md b/docs/user-guide/quick-start.md index b82530e7..292c6dd7 100644 --- a/docs/user-guide/quick-start.md +++ b/docs/user-guide/quick-start.md @@ -1,5 +1,9 @@ # 🚀 Quick Start +For construction, immutable parameter replacement, and differentiable +identification workflows, see +[Parameters, Updates, and Optimization](parameters-and-optimization.md). + **Get up and running with SoRoMoX in minutes!** This hands-on guide walks you through your first soft robot simulation with step-by-step examples. --- @@ -158,14 +162,14 @@ SoRoMoX uses an object-oriented design based on Equinox dataclasses. Systems are from soromox.systems import ( ArticulatedSoftRobot, ArticulatedSoftRobotParams, + LinkSpec, PCS, - PCSParams, ) # For spatial articulated soft robots robot = ArticulatedSoftRobot(params=ArticulatedSoftRobotParams(...)) # For soft continuum robots (PCS, GVS, HSA) -robot = PCS(params=PCSParams(...)) +robot = PCS.from_links([LinkSpec.circular(...)]) ``` **Key Benefits:** @@ -194,17 +198,20 @@ poses, gravity directions, and zero-gravity models. See for exact vectors and quaternions. ```python title="Parameter Structure" -params = PCSParams( - length=link_lengths, - radius=radii, - density=densities, - reference_strain=reference_strain, - young_modulus=young_modulus, - shear_modulus=shear_modulus, - material_damping_coefficient=material_damping_coefficient, -) -robot = PCS(params=params) -robot = robot.update_params(length=new_lengths) +links = [ + LinkSpec.circular( + length=length, + radius=radius, + density=density, + reference_strain=reference, + young_modulus=young, + shear_modulus=shear, + material_damping_coefficient=material_damping, + ) + for length, radius, density in zip(link_lengths, radii, densities) +] +robot = PCS.from_links(links) +robot = robot.update_link_params(length=new_lengths) ``` For routed tendons, pushing rods, muscles, and equivalent pressure chambers, @@ -248,25 +255,26 @@ Ready for something more advanced? Let's simulate a soft continuum robot: === "🌊 Continuum Robot" ```python - from soromox.systems import PlanarPCS, PlanarPCSParams + from soromox.systems import LinkSpec, PlanarPCS # Create a 3-segment soft robot num_segments = 3 - segment_lengths = 0.1 * jnp.ones((num_segments,)) - material_damping_coefficient = 318.0 - params = PlanarPCSParams( - length=segment_lengths, - radius=0.02 * jnp.ones((num_segments,)), - density=1070.0 * jnp.ones((num_segments,)), - reference_strain=jnp.tile(jnp.array([0.0, 1.0, 0.0]), num_segments), - young_modulus=2e3 * jnp.ones((num_segments,)), - shear_modulus=1e3 * jnp.ones((num_segments,)), - material_damping_coefficient=material_damping_coefficient, - ) + links = [ + LinkSpec.circular( + length=0.1, + radius=0.02, + density=1070.0, + reference_strain=[0.0, 1.0, 0.0], + young_modulus=2e3, + shear_modulus=1e3, + material_damping_coefficient=318.0, + ) + for _ in range(num_segments) + ] # Note: Damping helps stabilize simulations and represents material dissipation. # Initialize the PCS robot - robot = PlanarPCS(params=params) + robot = PlanarPCS.from_links(links) # Define configuration (strains) # Each segment has 3 strain components: [curvature, shear_x, shear_y] diff --git a/examples/control/actuation_space/setpoint_regulation_comparison.py b/examples/control/actuation_space/setpoint_regulation_comparison.py index 91492525..da09b01c 100644 --- a/examples/control/actuation_space/setpoint_regulation_comparison.py +++ b/examples/control/actuation_space/setpoint_regulation_comparison.py @@ -43,7 +43,7 @@ from soromox.rendering import Open3DRenderer from soromox.systems import ( PCS, - PCSParams, + LinkSpec, SystemState, ) @@ -69,18 +69,21 @@ def create_robot() -> tuple[PCS, int]: segment_lengths = 1e-1 * jnp.ones((num_segments,)) material_damping_coefficient = 362.0 - body_params = PCSParams( + body_params = PCS.params_from_links( + [ + LinkSpec.circular( + length=float(segment_lengths[index]), + radius=radius, + density=float(rho[index]), + young_modulus=2e3, + shear_modulus=1e3, + material_damping_coefficient=material_damping_coefficient, + reference_strain=[0, 0, 0, 1, 0, 0], + ) + for index in range(num_segments) + ], base_pose=jnp.array([0.5, 0.5, -0.5, 0.5, 0.0, 0.0, 0.0]), - length=segment_lengths, - radius=radius * jnp.ones((num_segments,)), - density=rho, gravity=jnp.array([0.0, 0.0, 9.81]), - young_modulus=2e3 * jnp.ones((num_segments,)), - shear_modulus=1e3 * jnp.ones((num_segments,)), - material_damping_coefficient=material_damping_coefficient, - reference_strain=jnp.tile( - jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), num_segments - ), ) # Tendon routing: 3 tendons at 120 degrees apart, parallel to backbone diff --git a/examples/control/operational_space/control_tendon_actuated_pcs_with_synergistic.py b/examples/control/operational_space/control_tendon_actuated_pcs_with_synergistic.py index 7d100faf..c242935c 100644 --- a/examples/control/operational_space/control_tendon_actuated_pcs_with_synergistic.py +++ b/examples/control/operational_space/control_tendon_actuated_pcs_with_synergistic.py @@ -32,7 +32,7 @@ from soromox.rendering import Open3DRenderer from soromox.systems import ( PCS, - PCSParams, + LinkSpec, SystemState, ) @@ -56,18 +56,21 @@ def main( segment_lengths = 1e-1 * jnp.ones((num_segments,)) material_damping_coefficient = 362.0 - body_params = PCSParams( + body_params = PCS.params_from_links( + [ + LinkSpec.circular( + length=float(segment_lengths[index]), + radius=2e-2, + density=float(rho[index]), + young_modulus=2e3, + shear_modulus=1e3, + material_damping_coefficient=material_damping_coefficient, + reference_strain=[0, 0, 0, 1, 0, 0], + ) + for index in range(num_segments) + ], base_pose=p0, - length=segment_lengths, - radius=2e-2 * jnp.ones((num_segments,)), - density=rho, gravity=jnp.array([0.0, 0.0, 9.81]), - young_modulus=2e3 * jnp.ones((num_segments,)), - shear_modulus=1e3 * jnp.ones((num_segments,)), - material_damping_coefficient=material_damping_coefficient, - reference_strain=jnp.tile( - jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), num_segments - ), ) # Tendons diff --git a/examples/simulation/gvs/simulate_gvs.py b/examples/simulation/gvs/simulate_gvs.py index 13e5ef97..8eb7f81a 100644 --- a/examples/simulation/gvs/simulate_gvs.py +++ b/examples/simulation/gvs/simulate_gvs.py @@ -15,10 +15,12 @@ ) from soromox.systems import ( GVS, - CrossSectionGeometry, + GVSSegment, + JointSpec, + LinkSpec, + StrainBasisSpec, SystemState, ) -from soromox.systems.gvs import GVSSegment, JointSpec, LinkSpec, StrainBasisSpec jnp.set_printoptions( threshold=jnp.inf, @@ -34,37 +36,34 @@ bases: list[StrainBasisSpec] = [] num_gauss_points: list[int] = [] - link1 = LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=1e6, - nu=0.5, - rho=1000, - eta=1e4, - L=0.3, - r_i=0.03, - r_f=0.03, + link1 = LinkSpec.circular( + young_modulus=1e6, + shear_modulus=1e6 / 3.0, + density=1000, + material_damping_coefficient=1e4, + length=0.3, + radius=0.03, + reference_strain=[0, 0, 0, 1, 0, 0], ) links.append(link1) joint1 = JointSpec(type="fixed") joints.append(joint1) basis1 = StrainBasisSpec( type="legendre", - active=[0, 1, 1, 0, 0, 0], - orders=[0, 0, 0, 0, 0, 0], - xi_ref=[0, 0, 0, 1, 0, 0], + strain_selector=[0, 1, 1, 0, 0, 0], + basis_order=[0, 0, 0, 0, 0, 0], ) bases.append(basis1) num_gauss_points.append(5) # Number of Gauss points for the first link - link2 = LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=1e6, - nu=0.5, - rho=1000, - eta=1e4, - L=0.3, - r_i=0.03, - r_f=0.03, + link2 = LinkSpec.circular( + young_modulus=1e6, + shear_modulus=1e6 / 3.0, + density=1000, + material_damping_coefficient=1e4, + length=0.3, + radius=0.03, + reference_strain=[0, 0, 0, 1, 0, 0], ) links.append(link2) joint2 = JointSpec(type="fixed") @@ -72,40 +71,12 @@ joints.append(joint2) basis2 = StrainBasisSpec( type="monomial", - active=[1, 1, 0, 0, 0, 0], - orders=[0, 0, 0, 0, 0, 0], - xi_ref=[0, 0, 0, 1, 0, 0], + strain_selector=[1, 1, 0, 0, 0, 0], + basis_order=[0, 0, 0, 0, 0, 0], ) bases.append(basis2) num_gauss_points.append(6) # Number of Gauss points for the second link - # link3 = LinkSpec( - # cross_section_geometry=CrossSectionGeometry.ELLIPTICAL, # Section type - # E=1e7, # Young's modulus in Pascals - # nu=0.4, # Poisson's ratio [-1, 0.5] - # rho=1050, # Density [kg/m^3] - # eta=1e4, # Damping coefficient - # L=0.3, # Length in meters - # a_i=0.04, # Initial semi-major axis in meters - # a_f=0.04, # Final semi-major axis in meters - # b_i=0.02, # Initial semi-minor axis in meters - # b_f=0.02 # Final semi-minor axis in meters - # ) - # links.append(link3) - # joint3 = JointSpec( - # type='revolute', # Prismatic joint - # axis='z', # Axis of translation - # ) - # joints.append(joint3) - # basis3 = StrainBasisSpec( - # type='chebyshev', # Type of basis - # active=[0, 1, 0, 1, 0, 0], # Degrees of freedom for each deformation type - # orders=[0, 0, 0, 0, 0, 0], # Order of basis functions for each deformation type - # xi_ref=[0, 0, 0, 1, 0, 0], # Reference strain values as vector - # ) - # bases.append(basis3) - # num_gauss_points.append(5) # Number of Gauss points for the third link - # ====================================================== # Robot initialization # ====================================================== diff --git a/examples/simulation/gvs/simulate_tendon_actuated_gvs.py b/examples/simulation/gvs/simulate_tendon_actuated_gvs.py index 0460b18e..b075a390 100644 --- a/examples/simulation/gvs/simulate_tendon_actuated_gvs.py +++ b/examples/simulation/gvs/simulate_tendon_actuated_gvs.py @@ -9,10 +9,13 @@ from soromox.rendering import Open3DRenderer from soromox.systems import ( GVS, - CrossSectionGeometry, + GVSSegment, + JointSpec, + LinearProfile, + LinkSpec, + StrainBasisSpec, SystemState, ) -from soromox.systems.gvs import GVSSegment, JointSpec, LinkSpec, StrainBasisSpec jax.config.update("jax_enable_x64", True) # jax.config.update("jax_platform_name", "gpu") # or "cpu" @@ -48,36 +51,38 @@ def statics_eq(q, args): # 2 link version # Link 1 -link1 = LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=3.04e5, - nu=0.45, - rho=1310.0, - eta=1e4, - L=0.0250 + 0.2550 + 0.0250, - r_i=0.01541, - r_f=0.00642, +link1 = LinkSpec.circular( + young_modulus=3.04e5, + shear_modulus=3.04e5 / 2.9, + density=1310.0, + material_damping_coefficient=1e4, + length=0.0250 + 0.2550 + 0.0250, + radius=LinearProfile(0.01541, 0.00642), + reference_strain=[0, 0, 0, 1, 0, 0], ) # Link 2 -link2 = LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=3.04e5, - nu=0.45, - rho=1310.0, - eta=1e4, - L=0.0550, - r_i=0.00642, - r_f=0.00480, +link2 = LinkSpec.circular( + young_modulus=3.04e5, + shear_modulus=3.04e5 / 2.9, + density=1310.0, + material_damping_coefficient=1e4, + length=0.0550, + radius=LinearProfile(0.00642, 0.00480), + reference_strain=[0, 0, 0, 1, 0, 0], ) joint1 = JointSpec(type="fixed") joint2 = JointSpec(type="fixed") basis1 = StrainBasisSpec( - type="monomial", active=[1, 1, 1, 1, 0, 0], orders=[1, 1, 1, 1, 0, 0] + type="monomial", + strain_selector=[1, 1, 1, 1, 0, 0], + basis_order=[1, 1, 1, 1, 0, 0], ) basis2 = StrainBasisSpec( - type="monomial", active=[0, 1, 1, 0, 0, 0], orders=[0, 0, 0, 0, 0, 0] + type="monomial", + strain_selector=[0, 1, 1, 0, 0, 0], + basis_order=[0, 0, 0, 0, 0, 0], ) diff --git a/examples/simulation/pcs/simulate_batched_tendon_actuated_pcs.py b/examples/simulation/pcs/simulate_batched_tendon_actuated_pcs.py index 5d17e3b9..6eefa6a0 100644 --- a/examples/simulation/pcs/simulate_batched_tendon_actuated_pcs.py +++ b/examples/simulation/pcs/simulate_batched_tendon_actuated_pcs.py @@ -24,7 +24,7 @@ ) from soromox.systems import ( PCS, - PCSParams, + LinkSpec, PCSStructure, SystemState, ) @@ -47,18 +47,21 @@ def build_robot() -> PCS: rho = 1070 * jnp.ones((num_segments,)) segment_lengths = 1e-1 * jnp.ones((num_segments,)) material_damping_coefficient = 362.0 - body_params = PCSParams( + body_params = PCS.params_from_links( + [ + LinkSpec.circular( + length=float(segment_lengths[index]), + radius=2e-2, + density=float(rho[index]), + young_modulus=2e3, + shear_modulus=1e3, + material_damping_coefficient=material_damping_coefficient, + reference_strain=[0, 0, 0, 1, 0, 0], + ) + for index in range(num_segments) + ], base_pose=jnp.array([0.5, 0.5, -0.5, 0.5, 0.0, 0.0, 0.0]), - length=segment_lengths, - radius=2e-2 * jnp.ones((num_segments,)), - density=rho, gravity=jnp.array([0.0, 0.0, 9.81]), - young_modulus=2e3 * jnp.ones((num_segments,)), - shear_modulus=1e3 * jnp.ones((num_segments,)), - material_damping_coefficient=material_damping_coefficient, - reference_strain=jnp.tile( - jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), num_segments - ), ) active_tendon_routing = ThreadlikeRouting.linear( intercept=jnp.array( diff --git a/examples/simulation/pcs/simulate_isupport.py b/examples/simulation/pcs/simulate_isupport.py index 876b1c1a..2145f3fb 100644 --- a/examples/simulation/pcs/simulate_isupport.py +++ b/examples/simulation/pcs/simulate_isupport.py @@ -18,7 +18,14 @@ ISupportVisualConfig, MatplotlibRenderer, ) -from soromox.systems import ISupport, ISupportParams, ISupportStructure, SystemState +from soromox.systems import ( + PCS, + ISupport, + ISupportParams, + ISupportStructure, + LinkSpec, + SystemState, +) if __name__ == "__main__": num_pneumatic_segments = 2 @@ -47,27 +54,37 @@ # Set this to None to divide each pneumatic segment equally according to # pcs_segment_counts. pcs_segment_lengths = jnp.array([95e-3, 95e-3, 60e-3, 60e-3, 60e-3]) - # Previous explicit damping used gamma_t = 806e-3 and gamma_r = 1.0e-3: - # D_i = L_i * diag([gamma_r, gamma_r, gamma_r, gamma_t, gamma_t, gamma_t]). - # The material damping coefficient below is the least-squares scalar c for - # D_i(c) = L_i * c * diag([Ix, 3Iy, 3Iz, 3A, A, A]), evaluated with - # I-SUPPORT's actuator cross-section geometry. Equivalently, - # c = _F / _F. - material_damping_coefficient = 1.96e3 + # Preserve the calibrated generalized damping used by the reference model. + gamma_rotational = 1.0e-3 + gamma_translational = 806e-3 + links = [ + LinkSpec.circular( + length=float(physical_segment_lengths[index]), + radius=float(physical_segment_radii[index]), + density=float(physical_segment_densities[index]), + young_modulus=E, + shear_modulus=G, + damping=float(physical_segment_lengths[index]) + * jnp.diag( + jnp.array( + [ + gamma_rotational, + gamma_rotational, + gamma_rotational, + gamma_translational, + gamma_translational, + gamma_translational, + ] + ) + ), + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + ) + for index in range(len(rigid_segment_selector)) + ] params = ISupportParams( base_pose=jnp.array([0.5, 0.5, 0.5, -0.5, 0.0, 0.0, 0.0]), - length=physical_segment_lengths, - radius=physical_segment_radii, - density=physical_segment_densities, gravity=jnp.array([0.0, 0.0, -9.81]), - young_modulus=E * jnp.ones((len(rigid_segment_selector),)), - shear_modulus=G * jnp.ones((len(rigid_segment_selector),)), - material_damping_coefficient=material_damping_coefficient - * jnp.ones((len(rigid_segment_selector),)), - reference_strain=jnp.tile( - jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), - len(rigid_segment_selector), - ), + link=PCS.params_from_links(links).link, chamber_inner_radius=6.39 * 1e-3 * jnp.ones((num_pneumatic_segments,)), chamber_outer_radius=7.79 * 1e-3 * jnp.ones((num_pneumatic_segments,)), chamber_distance=20 * 1e-3 * jnp.ones((num_pneumatic_segments,)), diff --git a/examples/simulation/pcs/simulate_pcs.py b/examples/simulation/pcs/simulate_pcs.py index 43a06b34..aa41e1b6 100644 --- a/examples/simulation/pcs/simulate_pcs.py +++ b/examples/simulation/pcs/simulate_pcs.py @@ -16,7 +16,7 @@ ViserRenderer, get_color_theme, ) -from soromox.systems import PCS, PCSParams, SystemState +from soromox.systems import PCS, LinkSpec, SystemState jnp.set_printoptions( threshold=jnp.inf, @@ -32,20 +32,23 @@ ) # Volumetric density of Dragon Skin 20 [kg/m^3] segment_lengths = 1e-1 * jnp.ones((num_segments,)) material_damping_coefficient = 362.0 - params = PCSParams( + params = PCS.params_from_links( + [ + LinkSpec.circular( + length=float(segment_lengths[index]), + radius=2e-2, + density=float(rho[index]), + young_modulus=2e3, + shear_modulus=1e3, + material_damping_coefficient=material_damping_coefficient, + reference_strain=[0, 0, 0, 1, 0, 0], + ) + for index in range(num_segments) + ], base_pose=jnp.array( [0.5, 0.5, -0.5, 0.5, 0.0, 0.0, 0.0] ), # Initial position and orientation - length=segment_lengths, - radius=2e-2 * jnp.ones((num_segments,)), - density=rho, gravity=jnp.array([0.0, 0.0, 9.81]), # Gravity vector [m/s^2] - young_modulus=2e3 * jnp.ones((num_segments,)), # Elastic modulus [Pa] - shear_modulus=1e3 * jnp.ones((num_segments,)), # Shear modulus [Pa] - material_damping_coefficient=material_damping_coefficient, - reference_strain=jnp.tile( - jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), num_segments - ), ) # ====================================================== diff --git a/examples/simulation/pcs/simulate_planar_pcs.py b/examples/simulation/pcs/simulate_planar_pcs.py index a71a43fa..f80156c6 100644 --- a/examples/simulation/pcs/simulate_planar_pcs.py +++ b/examples/simulation/pcs/simulate_planar_pcs.py @@ -9,7 +9,7 @@ jax.config.update("jax_enable_x64", True) # double precision from soromox.rendering import MatplotlibRenderer, OpenCVPlanarRenderer -from soromox.systems import PlanarPCS, PlanarPCSParams, SystemState +from soromox.systems import LinkSpec, PlanarPCS, SystemState jnp.set_printoptions( threshold=jnp.inf, @@ -32,16 +32,21 @@ ) # Volumetric density of Dragon Skin 20 [kg/m^3] segment_lengths = 1e-1 * jnp.ones((num_segments,)) material_damping_coefficient = 318.0 - params = PlanarPCSParams( + params = PlanarPCS.params_from_links( + [ + LinkSpec.circular( + length=float(segment_lengths[index]), + radius=2e-2, + density=float(rho[index]), + young_modulus=2e3, + shear_modulus=1e3, + material_damping_coefficient=material_damping_coefficient, + reference_strain=[0, 1, 0], + ) + for index in range(num_segments) + ], base_pose=jnp.array([jnp.pi / 2, 0.0, 0.0]), - length=segment_lengths, - radius=2e-2 * jnp.ones((num_segments,)), - density=rho, gravity=jnp.array([0.0, 9.81]), # gravity vector [m/s^2] UP! - young_modulus=2e3 * jnp.ones((num_segments,)), # Elastic modulus [Pa] - shear_modulus=1e3 * jnp.ones((num_segments,)), # Shear modulus [Pa] - material_damping_coefficient=material_damping_coefficient, - reference_strain=jnp.tile(jnp.array([0.0, 1.0, 0.0]), num_segments), ) # ====================================================== @@ -52,7 +57,7 @@ J, Jd = robot.jacobian_and_time_derivative( q=jnp.zeros((3 * num_segments,)), qd=jnp.zeros((3 * num_segments,)), - s=params.length[0], + s=params.link.length[0], ) # ===================================================== diff --git a/examples/simulation/pcs/simulate_tendon_actuated_pcs.py b/examples/simulation/pcs/simulate_tendon_actuated_pcs.py index 7a0df223..ad7a65ff 100644 --- a/examples/simulation/pcs/simulate_tendon_actuated_pcs.py +++ b/examples/simulation/pcs/simulate_tendon_actuated_pcs.py @@ -23,7 +23,7 @@ ) from soromox.systems import ( PCS, - PCSParams, + LinkSpec, SystemState, ) @@ -36,7 +36,9 @@ ) if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Simulate a tendon-actuated PCS robot.") + parser = argparse.ArgumentParser( + description="Simulate a tendon-actuated PCS robot." + ) parser.add_argument("--video-output", type=Path, default=DEFAULT_VIDEO_PATH) args = parser.parse_args() num_segments = 2 @@ -45,18 +47,21 @@ ) # Volumetric density of Dragon Skin 20 [kg/m^3] segment_lengths = 1e-1 * jnp.ones((num_segments,)) material_damping_coefficient = 362.0 - body_params = PCSParams( + body_params = PCS.params_from_links( + [ + LinkSpec.circular( + length=float(segment_lengths[index]), + radius=2e-2, + density=float(rho[index]), + young_modulus=2e3, + shear_modulus=1e3, + material_damping_coefficient=material_damping_coefficient, + reference_strain=[0, 0, 0, 1, 0, 0], + ) + for index in range(num_segments) + ], base_pose=jnp.array([0.5, 0.5, -0.5, 0.5, 0.0, 0.0, 0.0]), - length=segment_lengths, - radius=2e-2 * jnp.ones((num_segments,)), - density=rho, gravity=jnp.array([0.0, 0.0, 9.81]), - young_modulus=2e3 * jnp.ones((num_segments,)), - shear_modulus=1e3 * jnp.ones((num_segments,)), - material_damping_coefficient=material_damping_coefficient, - reference_strain=jnp.tile( - jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), num_segments - ), ) active_tendon_routing = ThreadlikeRouting.linear( intercept=2e-2 * jnp.array([[0.0, 1.0, 0.0], [0.0, -1.0, 0.0]]), diff --git a/examples/simulation/pcs/simulate_tendon_actuated_planar_pcs.py b/examples/simulation/pcs/simulate_tendon_actuated_planar_pcs.py index 7f25853c..24a86368 100644 --- a/examples/simulation/pcs/simulate_tendon_actuated_planar_pcs.py +++ b/examples/simulation/pcs/simulate_tendon_actuated_planar_pcs.py @@ -9,8 +9,8 @@ from soromox.actuation import ThreadlikeActuator, ThreadlikeRouting from soromox.rendering import MatplotlibRenderer from soromox.systems import ( + LinkSpec, PlanarPCS, - PlanarPCSParams, PlanarPCSStructure, SystemState, ) @@ -22,16 +22,21 @@ ) # Volumetric density of Dragon Skin 20 [kg/m^3] segment_lengths = 1e-1 * jnp.ones((num_segments,)) material_damping_coefficient = 318.0 - body = PlanarPCSParams( + body = PlanarPCS.params_from_links( + [ + LinkSpec.circular( + length=float(segment_lengths[index]), + radius=2e-2, + density=float(rho[index]), + young_modulus=5e3, + shear_modulus=1e3, + material_damping_coefficient=material_damping_coefficient, + reference_strain=[0, 1, 0], + ) + for index in range(num_segments) + ], base_pose=jnp.array([jnp.pi / 2, 0.0, 0.0]), - length=segment_lengths, - radius=2e-2 * jnp.ones((num_segments,)), - density=rho, gravity=0 * jnp.array([0.0, 9.81]), - young_modulus=5e3 * jnp.ones((num_segments,)), - shear_modulus=1e3 * jnp.ones((num_segments,)), - material_damping_coefficient=material_damping_coefficient, - reference_strain=jnp.tile(jnp.array([0.0, 1.0, 0.0]), num_segments), ) tendon_offsets = 2e-2 * jnp.array([[1.0, -1.0]]).repeat(num_segments, axis=0) active_tendon_routing = ThreadlikeRouting.linear( diff --git a/mkdocs.yml b/mkdocs.yml index 4db3fc07..5e853cfa 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -61,11 +61,13 @@ nav: - Installation: installation.md - User Guide: - Quick Start: user-guide/quick-start.md + - Parameters and Optimization: user-guide/parameters-and-optimization.md - Examples: user-guide/examples.md - API Reference: - Overview: api/overview.md - Systems: - Overview: api/systems/index.md + - Continuum Robot Components: api/systems/continuum-components.md - Articulated Systems: - Overview: api/systems/articulated/index.md - Pendulum: api/systems/pendulum/pendulum.md diff --git a/paper_results/secIVa_benchmarking_sequential_cpu/code/soromox/simulate_complex_gvs.py b/paper_results/secIVa_benchmarking_sequential_cpu/code/soromox/simulate_complex_gvs.py index d514fbd9..1f7e8ee8 100644 --- a/paper_results/secIVa_benchmarking_sequential_cpu/code/soromox/simulate_complex_gvs.py +++ b/paper_results/secIVa_benchmarking_sequential_cpu/code/soromox/simulate_complex_gvs.py @@ -15,10 +15,12 @@ ) from soromox.systems import ( GVS, - CrossSectionGeometry, + GVSSegment, + JointSpec, + LinkSpec, + StrainBasisSpec, SystemState, ) -from soromox.systems.gvs import GVSSegment, JointSpec, LinkSpec, StrainBasisSpec jax.config.update("jax_enable_x64", True) # jax.config.update("jax_platform_name", "gpu") @@ -41,47 +43,43 @@ bases: list[StrainBasisSpec] = [] num_gauss_points: list[int] = [] - link1 = LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=1e6, - nu=0.5, - rho=1000, - eta=1e4, - L=0.3, - r_i=0.03, - r_f=0.03, + link1 = LinkSpec.circular( + length=0.3, + radius=0.03, + density=1000, + young_modulus=1e6, + shear_modulus=1e6 / (2.0 * (1.0 + 0.5)), + material_damping_coefficient=1e4, + reference_strain=[0, 0, 0, 1, 0, 0], ) links.append(link1) - joint1 = JointSpec(type="Fixed") + joint1 = JointSpec.fixed() joints.append(joint1) basis1 = StrainBasisSpec( - type="Legendre", - active=[0, 1, 1, 1, 0, 0], - orders=[0, 1, 1, 1, 0, 0], - xi_ref=[0, 0, 0, 1, 0, 0], + type="legendre", + strain_selector=[0, 1, 1, 1, 0, 0], + basis_order=[0, 1, 1, 1, 0, 0], ) bases.append(basis1) num_gauss_points.append(5) # Number of Gauss points for the first link - link2 = LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=1e6, - nu=0.5, - rho=1000, - eta=1e4, - L=0.3, - r_i=0.03, - r_f=0.03, + link2 = LinkSpec.circular( + length=0.3, + radius=0.03, + density=1000, + young_modulus=1e6, + shear_modulus=1e6 / (2.0 * (1.0 + 0.5)), + material_damping_coefficient=1e4, + reference_strain=[0, 0, 0, 1, 0, 0], ) links.append(link2) - joint2 = JointSpec(type="Fixed") + joint2 = JointSpec.fixed() # joint2 = JointAttributes(jointtype='Revolute', axis='z') joints.append(joint2) basis2 = StrainBasisSpec( - type="Legendre", - active=[0, 1, 1, 1, 0, 0], - orders=[0, 1, 1, 1, 0, 0], - xi_ref=[0, 0, 0, 1, 0, 0], + type="legendre", + strain_selector=[0, 1, 1, 1, 0, 0], + basis_order=[0, 1, 1, 1, 0, 0], ) bases.append(basis2) num_gauss_points.append(5) # Number of Gauss points for the second link diff --git a/paper_results/secIVa_benchmarking_sequential_cpu/code/soromox/simulate_planar_pcs.py b/paper_results/secIVa_benchmarking_sequential_cpu/code/soromox/simulate_planar_pcs.py index 31f9d0a3..a3e90d49 100644 --- a/paper_results/secIVa_benchmarking_sequential_cpu/code/soromox/simulate_planar_pcs.py +++ b/paper_results/secIVa_benchmarking_sequential_cpu/code/soromox/simulate_planar_pcs.py @@ -9,7 +9,7 @@ from diffrax import Tsit5 from soromox.rendering import MatplotlibRenderer, OpenCVPlanarRenderer -from soromox.systems import PlanarPCS, PlanarPCSParams, SystemState +from soromox.systems import LinkSpec, PlanarPCS, SystemState jax.config.update("jax_enable_x64", True) # double precision @@ -41,16 +41,25 @@ def _planar_position_from_pose(chi_ts: jax.Array) -> jax.Array: * segment_lengths[:, None] ).flatten() ) - params = PlanarPCSParams( + links = [ + LinkSpec.circular( + length=float(segment_lengths[index]), + radius=3e-2, + density=float(rho[index]), + young_modulus=1e6, + shear_modulus=3.333333e5, + damping=damping_matrix[ + 3 * index : 3 * (index + 1), + 3 * index : 3 * (index + 1), + ], + reference_strain=[0.0, 1.0, 0.0], + ) + for index in range(num_segments) + ] + params = PlanarPCS.params_from_links( + links, base_pose=jnp.array([jnp.pi / 2, 0.0, 0.0]), - length=segment_lengths, - radius=3e-2 * jnp.ones((num_segments,)), - density=rho, gravity=jnp.array([-9.81 * 1, 9.81 * 0]), # gravity vector [m/s^2] UP! - young_modulus=1e6 * jnp.ones((num_segments,)), # Elastic modulus [Pa] - shear_modulus=3.333333e5 * jnp.ones((num_segments,)), # Shear modulus [Pa] - damping_matrix=damping_matrix, - reference_strain=jnp.tile(jnp.array([0.0, 1.0, 0.0]), num_segments), ) # ====================================================== @@ -61,7 +70,7 @@ def _planar_position_from_pose(chi_ts: jax.Array) -> jax.Array: J, Jd = robot.jacobian_and_time_derivative( q=jnp.zeros((3 * num_segments,)), qd=jnp.zeros((3 * num_segments,)), - s=params.length[0], + s=params.link.length[0], ) # ===================================================== diff --git a/paper_results/secIVa_benchmarking_sequential_cpu/code/soromox/simulate_spatial_pcs.py b/paper_results/secIVa_benchmarking_sequential_cpu/code/soromox/simulate_spatial_pcs.py index 768954f4..b153b0bb 100644 --- a/paper_results/secIVa_benchmarking_sequential_cpu/code/soromox/simulate_spatial_pcs.py +++ b/paper_results/secIVa_benchmarking_sequential_cpu/code/soromox/simulate_spatial_pcs.py @@ -8,7 +8,7 @@ from diffrax import Tsit5 from soromox.rendering import MatplotlibRenderer -from soromox.systems import PCS, PCSParams, SystemState +from soromox.systems import PCS, LinkSpec, SystemState jax.config.update("jax_enable_x64", True) # double precision @@ -34,20 +34,27 @@ * segment_lengths[:, None] ).flatten() ) - params = PCSParams( + links = [ + LinkSpec.circular( + length=float(segment_lengths[index]), + radius=3e-2, + density=float(rho[index]), + young_modulus=1e6, + shear_modulus=3.333333e5, + damping=damping_matrix[ + 6 * index : 6 * (index + 1), + 6 * index : 6 * (index + 1), + ], + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + ) + for index in range(num_segments) + ] + params = PCS.params_from_links( + links, base_pose=jnp.array( [0.5, 0.5, -0.5, 0.5, 0.0, 0.0, 0.0] ), # Initial position and orientation - length=segment_lengths, - radius=3e-2 * jnp.ones((num_segments,)), - density=rho, gravity=jnp.array([-9.81 * 1, -9.81 * 1, 9.81 * 0]), # Gravity vector [m/s^2] - young_modulus=1e6 * jnp.ones((num_segments,)), # Elastic modulus [Pa] - shear_modulus=3.333333e5 * jnp.ones((num_segments,)), # Shear modulus [Pa] - damping_matrix=damping_matrix, - reference_strain=jnp.tile( - jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), num_segments - ), ) # ====================================================== diff --git a/paper_results/secIVa_benchmarking_sequential_cpu/code/soromox/simulate_tendon_driven_gvs.py b/paper_results/secIVa_benchmarking_sequential_cpu/code/soromox/simulate_tendon_driven_gvs.py index 84296174..04313e26 100644 --- a/paper_results/secIVa_benchmarking_sequential_cpu/code/soromox/simulate_tendon_driven_gvs.py +++ b/paper_results/secIVa_benchmarking_sequential_cpu/code/soromox/simulate_tendon_driven_gvs.py @@ -11,10 +11,12 @@ from soromox.actuation import ThreadlikeActuator, ThreadlikeRouting from soromox.systems import ( GVS, - CrossSectionGeometry, + GVSSegment, + JointSpec, + LinkSpec, + StrainBasisSpec, SystemState, ) -from soromox.systems.gvs import GVSSegment, JointSpec, LinkSpec, StrainBasisSpec jax.config.update("jax_enable_x64", True) # jax.config.update("jax_platform_name", "gpu") # or "cpu" @@ -51,36 +53,38 @@ def statics_eq(q, args): # 2 link version # Link 1 -link1 = LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=1e6, - nu=0.5, - rho=1000.0, - eta=1e4, - L=0.3, - r_i=0.03, - r_f=0.03, +link1 = LinkSpec.circular( + length=0.3, + radius=0.03, + density=1000.0, + young_modulus=1e6, + shear_modulus=1e6 / (2.0 * (1.0 + 0.5)), + material_damping_coefficient=1e4, + reference_strain=[0, 0, 0, 1, 0, 0], ) # Link 2 -link2 = LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=1e6, - nu=0.5, - rho=1000.0, - eta=1e4, - L=0.3, - r_i=0.03, - r_f=0.03, +link2 = LinkSpec.circular( + length=0.3, + radius=0.03, + density=1000.0, + young_modulus=1e6, + shear_modulus=1e6 / (2.0 * (1.0 + 0.5)), + material_damping_coefficient=1e4, + reference_strain=[0, 0, 0, 1, 0, 0], ) joint1 = JointSpec(type="Fixed") joint2 = JointSpec(type="Fixed") basis1 = StrainBasisSpec( - type="Legendre", active=[1, 1, 1, 1, 0, 0], orders=[1, 1, 1, 1, 0, 0] + type="legendre", + strain_selector=[1, 1, 1, 1, 0, 0], + basis_order=[1, 1, 1, 1, 0, 0], ) basis2 = StrainBasisSpec( - type="Legendre", active=[1, 1, 1, 1, 0, 0], orders=[1, 1, 1, 1, 0, 0] + type="legendre", + strain_selector=[1, 1, 1, 1, 0, 0], + basis_order=[1, 1, 1, 1, 0, 0], ) diff --git a/paper_results/secVa_system_identification/code/identify_soft_tentacle_parameters.py b/paper_results/secVa_system_identification/code/identify_soft_tentacle_parameters.py index 00bfeab8..b77b01ca 100644 --- a/paper_results/secVa_system_identification/code/identify_soft_tentacle_parameters.py +++ b/paper_results/secVa_system_identification/code/identify_soft_tentacle_parameters.py @@ -14,9 +14,12 @@ from soromox.actuation import ThreadlikeActuator, ThreadlikeRouting from soromox.systems import ( GVS, - CrossSectionGeometry, + GVSSegment, + JointSpec, + LinearProfile, + LinkSpec, + StrainBasisSpec, ) -from soromox.systems.gvs import GVSSegment, JointSpec, LinkSpec, StrainBasisSpec jax.config.update("jax_enable_x64", True) @@ -453,35 +456,37 @@ def loss_fn(params): # 2 link version # Link 1 - link1 = LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=5.05e5, - nu=0.45, - rho=1500.0, - eta=1e4, - L=0.0250 + 0.2550 + 0.0250, - r_i=0.01541, - r_f=0.00642, + link1 = LinkSpec.circular( + length=0.0250 + 0.2550 + 0.0250, + radius=LinearProfile(base=0.01541, tip=0.00642), + density=1500.0, + young_modulus=5.05e5, + shear_modulus=5.05e5 / (2.0 * (1.0 + 0.45)), + material_damping_coefficient=1e4, + reference_strain=[0, 0, 0, 1, 0, 0], ) # Link 2 - link2 = LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=5.05e5, - nu=0.45, - rho=1500.0, - eta=1e4, - L=0.0550, - r_i=0.00642, - r_f=0.00480, + link2 = LinkSpec.circular( + length=0.0550, + radius=LinearProfile(base=0.00642, tip=0.00480), + density=1500.0, + young_modulus=5.05e5, + shear_modulus=5.05e5 / (2.0 * (1.0 + 0.45)), + material_damping_coefficient=1e4, + reference_strain=[0, 0, 0, 1, 0, 0], ) joint1 = JointSpec(type="fixed") joint2 = JointSpec(type="fixed") basis1 = StrainBasisSpec( - type="monomial", active=[1, 1, 1, 1, 0, 0], orders=[0, 1, 1, 1, 0, 0] + type="monomial", + strain_selector=[1, 1, 1, 1, 0, 0], + basis_order=[0, 1, 1, 1, 0, 0], ) basis2 = StrainBasisSpec( - type="monomial", active=[0, 1, 1, 0, 0, 0], orders=[0, 0, 0, 0, 0, 0] + type="monomial", + strain_selector=[0, 1, 1, 0, 0, 0], + basis_order=[0, 0, 0, 0, 0, 0], ) num_gauss_points = [8, 8] @@ -532,9 +537,9 @@ def loss_fn(params): n_links = int(robot.num_segments) # E and nu initial values (same for all links) - E0 = float(link1.E) - nu0 = float(link1.nu) - rho0 = float(link1.rho) + E0 = float(link1.young_modulus) + nu0 = 0.45 + rho0 = float(link1.density) E_init = jnp.asarray(E0, dtype=jnp.float64) # shape: () nu_init = jnp.asarray(nu0, dtype=jnp.float64) # shape: () rho_init = jnp.asarray(rho0, dtype=jnp.float64) # shape: () @@ -754,25 +759,23 @@ def marker_rmse(errors): finish_figure(show=not args.no_show) # Markers comparison visualization - link1_hat = LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=float(E_hat), - nu=float(nu_hat), - rho=float(rho_hat), - eta=1e4, - L=0.0250 + 0.2550 + 0.0250, - r_i=0.01541, - r_f=0.00642, + link1_hat = LinkSpec.circular( + length=0.0250 + 0.2550 + 0.0250, + radius=LinearProfile(base=0.01541, tip=0.00642), + density=float(rho_hat), + young_modulus=float(E_hat), + shear_modulus=float(E_hat / (2.0 * (1.0 + nu_hat))), + material_damping_coefficient=1e4, + reference_strain=[0, 0, 0, 1, 0, 0], ) - link2_hat = LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=float(E_hat), - nu=float(nu_hat), - rho=float(rho_hat), - eta=1e4, - L=0.0550, - r_i=0.00642, - r_f=0.00480, + link2_hat = LinkSpec.circular( + length=0.0550, + radius=LinearProfile(base=0.00642, tip=0.00480), + density=float(rho_hat), + young_modulus=float(E_hat), + shear_modulus=float(E_hat / (2.0 * (1.0 + nu_hat))), + material_damping_coefficient=1e4, + reference_strain=[0, 0, 0, 1, 0, 0], ) robot_hat = GVS.from_segments( [ diff --git a/paper_results/secVa_system_identification/code/identify_soft_tentacle_residual.py b/paper_results/secVa_system_identification/code/identify_soft_tentacle_residual.py index 3b03bb36..4fb67382 100644 --- a/paper_results/secVa_system_identification/code/identify_soft_tentacle_residual.py +++ b/paper_results/secVa_system_identification/code/identify_soft_tentacle_residual.py @@ -15,9 +15,9 @@ from soromox.actuation import ThreadlikeActuator, ThreadlikeRouting from soromox.systems import ( GVS, - CrossSectionGeometry, GVSSegment, JointSpec, + LinearProfile, LinkSpec, StrainBasisSpec, ) @@ -61,6 +61,7 @@ def finish_figure(*, show: bool) -> None: else: plt.close() + jax.config.update("jax_enable_x64", True) # print("JAX default backend:", jax.default_backend()) # print("JAX devices:", jax.devices()) @@ -792,35 +793,37 @@ def loss_model(m: TauNN): ### BODY DEFINITION OF THE SOFT ROBOT ### # Link 1 - link1 = LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=3.15e5, - nu=0.45, - rho=1320.0, - eta=1e4, - L=0.0250 + 0.2550 + 0.0250, - r_i=0.01541, - r_f=0.00642, + link1 = LinkSpec.circular( + length=0.0250 + 0.2550 + 0.0250, + radius=LinearProfile(base=0.01541, tip=0.00642), + density=1320.0, + young_modulus=3.15e5, + shear_modulus=3.15e5 / (2.0 * (1.0 + 0.45)), + material_damping_coefficient=1e4, + reference_strain=[0, 0, 0, 1, 0, 0], ) # Link 2 - link2 = LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=3.15e5, - nu=0.45, - rho=1320.0, - eta=1e4, - L=0.0550, - r_i=0.00642, - r_f=0.00480, + link2 = LinkSpec.circular( + length=0.0550, + radius=LinearProfile(base=0.00642, tip=0.00480), + density=1320.0, + young_modulus=3.15e5, + shear_modulus=3.15e5 / (2.0 * (1.0 + 0.45)), + material_damping_coefficient=1e4, + reference_strain=[0, 0, 0, 1, 0, 0], ) joint1 = JointSpec(type="fixed") joint2 = JointSpec(type="fixed") basis1 = StrainBasisSpec( - type="monomial", active=[1, 1, 1, 1, 0, 0], orders=[0, 1, 1, 1, 0, 0] + type="monomial", + strain_selector=[1, 1, 1, 1, 0, 0], + basis_order=[0, 1, 1, 1, 0, 0], ) basis2 = StrainBasisSpec( - type="monomial", active=[0, 1, 1, 0, 0, 0], orders=[0, 0, 0, 0, 0, 0] + type="monomial", + strain_selector=[0, 1, 1, 0, 0, 0], + basis_order=[0, 0, 0, 0, 0, 0], ) num_gauss_points = [8, 8] diff --git a/paper_results/secVc_model_based_control/configuration_space_comparison/code/configuration_space_comparison_simulation.py b/paper_results/secVc_model_based_control/configuration_space_comparison/code/configuration_space_comparison_simulation.py index 860943b3..24a09254 100644 --- a/paper_results/secVc_model_based_control/configuration_space_comparison/code/configuration_space_comparison_simulation.py +++ b/paper_results/secVc_model_based_control/configuration_space_comparison/code/configuration_space_comparison_simulation.py @@ -38,7 +38,7 @@ PotentialCancellationRegulator, PotentialCompensationRegulator, ) -from soromox.systems import PCS, PCSParams, SystemState +from soromox.systems import PCS, LinkSpec, SystemState CASE_DIR = Path(__file__).parent.parent DATA_DIR = CASE_DIR / "data" @@ -191,22 +191,25 @@ def create_robot( * segment_lengths[:, None] ).flatten() ) - params = PCSParams( + links = [ + LinkSpec.circular( + length=float(segment_lengths[index]), + radius=2e-2, + density=float(rho[index]), + young_modulus=2e3, + shear_modulus=1e3, + damping=damping_matrix[ + 6 * index : 6 * (index + 1), 6 * index : 6 * (index + 1) + ], + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + ) + for index in range(num_segments) + ] + robot = PCS.from_links( + links, base_pose=jnp.asarray(base_pose), - length=segment_lengths, - radius=2e-2 * jnp.ones((num_segments,)), - density=rho, gravity=jnp.array([0.0, 0.0, -9.81]), - young_modulus=2e3 * jnp.ones((num_segments,)), - shear_modulus=1e3 * jnp.ones((num_segments,)), - damping_matrix=damping_matrix, - reference_strain=jnp.tile( - jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), - num_segments, - ), ) - - robot = PCS(params=params) return robot, int(robot.num_active_strains), num_segments diff --git a/paper_results/secVc_model_based_control/operational_space_impedance_control/code/compare_impedance_feedback_linearization.py b/paper_results/secVc_model_based_control/operational_space_impedance_control/code/compare_impedance_feedback_linearization.py index ea4ac5ec..1893effb 100644 --- a/paper_results/secVc_model_based_control/operational_space_impedance_control/code/compare_impedance_feedback_linearization.py +++ b/paper_results/secVc_model_based_control/operational_space_impedance_control/code/compare_impedance_feedback_linearization.py @@ -17,7 +17,7 @@ ReferenceTrajectory, ) from soromox.coordinate_transformations import OperationalSpaceDynamics -from soromox.systems import PCS, PCSParams, SystemState +from soromox.systems import PCS, LinkSpec, SystemState from soromox.utils.geometry.rotations import ( rotation_matrix_to_rotation_vector, rotation_vector_to_rotation_matrix, @@ -33,21 +33,22 @@ def build_robot() -> PCS: """Construct the fully actuated two-segment PCS benchmark robot.""" num_segments = 2 - params = PCSParams( + return PCS.from_links( + [ + LinkSpec.circular( + length=0.1, + radius=0.02, + density=1070.0, + young_modulus=2e3, + shear_modulus=1e3, + material_damping_coefficient=362.0, + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + ) + for _ in range(num_segments) + ], base_pose=jnp.array([0.5, 0.5, -0.5, 0.5, 0.0, 0.0, 0.0]), - length=0.1 * jnp.ones((num_segments,)), - radius=0.02 * jnp.ones((num_segments,)), - density=1070.0 * jnp.ones((num_segments,)), gravity=jnp.array([0.0, 0.0, 9.81]), - young_modulus=2e3 * jnp.ones((num_segments,)), - shear_modulus=1e3 * jnp.ones((num_segments,)), - material_damping_coefficient=362.0, - reference_strain=jnp.tile( - jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), - num_segments, - ), ) - return PCS(params=params) def minimum_jerk_ramp(t: jax.Array, duration: float = 0.25) -> jax.Array: diff --git a/paper_results/secVc_model_based_control/operational_space_impedance_control/code/operational_space_impedance_common.py b/paper_results/secVc_model_based_control/operational_space_impedance_control/code/operational_space_impedance_common.py index 027dabf2..1dc3ec6f 100644 --- a/paper_results/secVc_model_based_control/operational_space_impedance_control/code/operational_space_impedance_common.py +++ b/paper_results/secVc_model_based_control/operational_space_impedance_control/code/operational_space_impedance_common.py @@ -27,7 +27,7 @@ from soromox.control import OperationalSpaceImpedanceControlTracker from soromox.coordinate_transformations import OperationalSpaceDynamics -from soromox.systems import PCS, PCSParams, SystemState +from soromox.systems import PCS, LinkSpec, PCSParams, SystemState jax.config.update("jax_enable_x64", True) @@ -113,21 +113,26 @@ def build_problem(trajectory_config: TaskSpaceTrajectoryConfig) -> PCSImpedanceP * segment_lengths[:, None] ).flatten() ) - params = PCSParams( + links = [ + LinkSpec.circular( + length=float(segment_lengths[index]), + radius=2e-2, + density=float(rho[index]), + young_modulus=2e3, + shear_modulus=1e3, + damping=damping_matrix[ + 6 * index : 6 * (index + 1), 6 * index : 6 * (index + 1) + ], + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + ) + for index in range(num_segments) + ] + robot = PCS.from_links( + links, base_pose=p0, - length=segment_lengths, - radius=2e-2 * jnp.ones((num_segments,)), - density=rho, gravity=jnp.array([0.0, 0.0, -9.81]), - young_modulus=2e3 * jnp.ones((num_segments,)), - shear_modulus=1e3 * jnp.ones((num_segments,)), - damping_matrix=damping_matrix, - reference_strain=jnp.tile( - jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), num_segments - ), ) - - robot = PCS(params=params) + params = robot.params num_dofs = robot.num_active_strains total_length = float(jnp.sum(segment_lengths)) q0 = jnp.zeros((num_dofs,)) diff --git a/paper_results/secVd_control_gain_optimization/code/control_gain_optimization_with_collocated.py b/paper_results/secVd_control_gain_optimization/code/control_gain_optimization_with_collocated.py index 18a79030..cae52c7a 100644 --- a/paper_results/secVd_control_gain_optimization/code/control_gain_optimization_with_collocated.py +++ b/paper_results/secVd_control_gain_optimization/code/control_gain_optimization_with_collocated.py @@ -18,7 +18,7 @@ from soromox.rendering import Open3DRenderer from soromox.systems import ( PCS, - PCSParams, + LinkSpec, SystemState, ) @@ -124,16 +124,24 @@ def evaluate_closed_loop_system( * segment_lengths[:, None] ).flatten() ) -body_params = PCSParams( +body_params = PCS.params_from_links( + [ + LinkSpec.circular( + length=float(segment_lengths[index]), + radius=radius, + density=float(rho[index]), + young_modulus=2e3, + shear_modulus=1e3, + damping=damping_matrix[ + 6 * index : 6 * (index + 1), + 6 * index : 6 * (index + 1), + ], + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + ) + for index in range(num_segments) + ], base_pose=jnp.array([0.5, 0.5, -0.5, 0.5, 0.0, 0.0, 0.0]), - length=segment_lengths, - radius=radius * jnp.ones((num_segments,)), - density=rho, gravity=jnp.array([0.0, 0.0, 9.81]), - young_modulus=2e3 * jnp.ones((num_segments,)), - shear_modulus=1e3 * jnp.ones((num_segments,)), - damping_matrix=damping_matrix, - reference_strain=jnp.tile(jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), num_segments), ) # Tendon routing: 3 tendons at 120 degrees apart, parallel to backbone diff --git a/paper_results/secVd_control_gain_optimization/code/control_gain_optimization_with_synergistic.py b/paper_results/secVd_control_gain_optimization/code/control_gain_optimization_with_synergistic.py index 943262ee..67c159f3 100644 --- a/paper_results/secVd_control_gain_optimization/code/control_gain_optimization_with_synergistic.py +++ b/paper_results/secVd_control_gain_optimization/code/control_gain_optimization_with_synergistic.py @@ -22,7 +22,7 @@ from soromox.rendering import Open3DRenderer from soromox.systems import ( PCS, - PCSParams, + LinkSpec, SystemState, ) @@ -133,16 +133,24 @@ def evaluate_closed_loop_system( * segment_lengths[:, None] ).flatten() ) -body_params = PCSParams( +body_params = PCS.params_from_links( + [ + LinkSpec.circular( + length=float(segment_lengths[index]), + radius=radius, + density=float(rho[index]), + young_modulus=2e3, + shear_modulus=1e3, + damping=damping_matrix[ + 6 * index : 6 * (index + 1), + 6 * index : 6 * (index + 1), + ], + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + ) + for index in range(num_segments) + ], base_pose=jnp.array([0.5, 0.5, -0.5, 0.5, 0.0, 0.0, 0.0]), - length=segment_lengths, - radius=radius * jnp.ones((num_segments,)), - density=rho, gravity=jnp.array([0.0, 0.0, 9.81]), - young_modulus=2e3 * jnp.ones((num_segments,)), - shear_modulus=1e3 * jnp.ones((num_segments,)), - damping_matrix=damping_matrix, - reference_strain=jnp.tile(jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), num_segments), ) # Tendon routing: 3 tendons at 120 degrees apart, parallel to backbone diff --git a/paper_results/secVe_safety_constrained_control/code/pcs_cf_cbf_clf_common.py b/paper_results/secVe_safety_constrained_control/code/pcs_cf_cbf_clf_common.py index a4c272ba..95da9006 100644 --- a/paper_results/secVe_safety_constrained_control/code/pcs_cf_cbf_clf_common.py +++ b/paper_results/secVe_safety_constrained_control/code/pcs_cf_cbf_clf_common.py @@ -14,7 +14,7 @@ from soromox.actuation import ThreadlikeActuator, ThreadlikeRouting from soromox.systems import ( PCS, - PCSParams, + LinkSpec, SystemState, ) @@ -295,18 +295,24 @@ def build_simulation_setup() -> SimulationSetup: .flatten() ) - body_params = PCSParams( + body_params = PCS.params_from_links( + [ + LinkSpec.circular( + length=float(segment_length[index]), + radius=float(backbone_radius[index]), + density=float(rho[index]), + young_modulus=20e3, + shear_modulus=20e3, + damping=damping_matrix[ + 6 * index : 6 * (index + 1), + 6 * index : 6 * (index + 1), + ], + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + ) + for index in range(num_segments) + ], base_pose=jnp.array([0.5, 0.5, -0.5, 0.5, 0.0, 0.0, 0.0]), - length=segment_length, - radius=backbone_radius, - density=rho, gravity=jnp.array([0.0, 0.0, 9.81]), - young_modulus=20e3 * jnp.ones((num_segments,)), - shear_modulus=20e3 * jnp.ones((num_segments,)), - damping_matrix=damping_matrix, - reference_strain=jnp.tile( - jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), num_segments - ), ) active_tendon_routing = ThreadlikeRouting.linear( @@ -354,9 +360,9 @@ def build_simulation_setup() -> SimulationSetup: "strain_selector": strain_selector, "base_pose": body_params.base_pose, "gravity": body_params.gravity, - "young_modulus": body_params.young_modulus, - "shear_modulus": body_params.shear_modulus, - "reference_strain": body_params.reference_strain, + "young_modulus": 20e3 * jnp.ones((num_segments,)), + "shear_modulus": 20e3 * jnp.ones((num_segments,)), + "reference_strain": body_params.link.reference_strain, "obs_centers": jnp.array( [ [0.10, 0.08, 0.24], diff --git a/paper_results/secVf_parallel_rl/code/parallel_soromox_env.py b/paper_results/secVf_parallel_rl/code/parallel_soromox_env.py index 5e466d33..b04b550c 100644 --- a/paper_results/secVf_parallel_rl/code/parallel_soromox_env.py +++ b/paper_results/secVf_parallel_rl/code/parallel_soromox_env.py @@ -20,10 +20,7 @@ from stable_baselines3.common.vec_env import VecEnv from soromox.actuation import ThreadlikeActuator, ThreadlikeRouting -from soromox.systems.pcs import ( - PCS, - PCSParams, -) +from soromox.systems import PCS, LinkSpec from soromox.systems.system_state import SystemState jax.config.update("jax_enable_x64", True) @@ -113,18 +110,24 @@ def build_arm( ).flatten() ) - body_params = PCSParams( + body_params = PCS.params_from_links( + [ + LinkSpec.circular( + length=float(segment_lengths[index]), + radius=radius, + density=density, + young_modulus=youngs_modulus, + shear_modulus=shear_modulus, + damping=damping_matrix[ + 6 * index : 6 * (index + 1), + 6 * index : 6 * (index + 1), + ], + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + ) + for index in range(num_segments) + ], base_pose=jnp.array([0.5, 0.5, -0.5, 0.5, 0.0, 0.0, 0.0]), - length=segment_lengths, - radius=jnp.ones((num_segments,)) * radius, - density=jnp.ones((num_segments,)) * density, gravity=jnp.array([0.0, 0.0, 0.0]), - young_modulus=jnp.ones((num_segments,)) * youngs_modulus, - shear_modulus=jnp.ones((num_segments,)) * shear_modulus, - damping_matrix=damping_matrix, - reference_strain=jnp.tile( - jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), num_segments - ), ) active_tendon_routing = ThreadlikeRouting.linear( intercept=jnp.array( diff --git a/paper_results/secVf_parallel_rl/code/render_rl_video.py b/paper_results/secVf_parallel_rl/code/render_rl_video.py index fce0fe05..9fcf1f02 100644 --- a/paper_results/secVf_parallel_rl/code/render_rl_video.py +++ b/paper_results/secVf_parallel_rl/code/render_rl_video.py @@ -15,7 +15,7 @@ from soromox.rendering.color_config import BackboneColorConfig from soromox.systems import ( PCS, - PCSParams, + LinkSpec, ) SCRIPT_DIR = Path(__file__).resolve().parent.parent @@ -70,16 +70,24 @@ def build_rl_robot(arm_length: float = 0.25, arm_radius: float = 0.025) -> PCS: "idx_seg_att": jnp.zeros(4, dtype=jnp.int32), } - body_params = PCSParams( + body_params = PCS.params_from_links( + [ + LinkSpec.circular( + length=float(segment_length[index]), + radius=float(backbone_radius[index]), + density=float(rho[index]), + young_modulus=20e3, + shear_modulus=20e3, + damping=damping_matrix[ + 6 * index : 6 * (index + 1), + 6 * index : 6 * (index + 1), + ], + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + ) + for index in range(num_segments) + ], base_pose=jnp.array([0.5, 0.5, -0.5, 0.5, 0.0, 0.0, 0.0]), - length=segment_length, - radius=backbone_radius, - density=rho, gravity=jnp.array([0.0, 0.0, 9.81]), - young_modulus=20e3 * jnp.ones((num_segments,)), - shear_modulus=20e3 * jnp.ones((num_segments,)), - damping_matrix=damping_matrix, - reference_strain=jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), ) active_tendon_routing = ThreadlikeRouting.linear( diff --git a/src/soromox/rendering/open3d_renderer.py b/src/soromox/rendering/open3d_renderer.py index 2b0ca7a8..1819ecbe 100644 --- a/src/soromox/rendering/open3d_renderer.py +++ b/src/soromox/rendering/open3d_renderer.py @@ -46,7 +46,8 @@ from soromox.rendering.camera_config import CameraConfig from soromox.rendering.color_config import RendererColorConfig, ensure_rgba from soromox.rendering.video_encoding import FFmpegVideoWriter, VideoEncodingConfig -from soromox.systems.soft_robot import CrossSectionGeometry, SoftRobot +from soromox.systems.components import CrossSectionGeometry +from soromox.systems.soft_robot import SoftRobot # ====================================================================================== # Geometry helper functions (module-level, stateless) diff --git a/src/soromox/rendering/opencv_planar_renderer.py b/src/soromox/rendering/opencv_planar_renderer.py index 159dde2d..ac5df78b 100644 --- a/src/soromox/rendering/opencv_planar_renderer.py +++ b/src/soromox/rendering/opencv_planar_renderer.py @@ -11,7 +11,8 @@ from jax import Array from soromox.rendering.opencv_base import BaseOpenCVRenderer -from soromox.systems.soft_robot import CrossSectionGeometry, SoftRobot +from soromox.systems.components import CrossSectionGeometry +from soromox.systems.soft_robot import SoftRobot class OpenCVPlanarRenderer(BaseOpenCVRenderer): diff --git a/src/soromox/systems/__init__.py b/src/soromox/systems/__init__.py index 00bd64c5..b8ec9eea 100644 --- a/src/soromox/systems/__init__.py +++ b/src/soromox/systems/__init__.py @@ -7,8 +7,19 @@ ArticulatedSoftRobotParams, McKibbenActuatedUMArmParams, ) +from soromox.systems.components import ( + ContinuumLinkParams, + CrossSectionGeometry, + CrossSectionParams, + IsotropicMaterialParams, + JointParams, + JointSpec, + LinearProfile, + LinkSpec, + shear_modulus_from_poisson_ratio, +) from soromox.systems.dynamical_system import DynamicalSystem -from soromox.systems.gvs.params import GVSLinkParams, GVSParams +from soromox.systems.gvs.params import GVSParams from soromox.systems.gvs.structures import ( GVSJointStructure, GVSLinkStructure, @@ -36,18 +47,13 @@ PlanarPCSStructure, ) from soromox.systems.pendulum.params import PendulumParams -from soromox.systems.soft_robot import ( - CrossSectionGeometry, - SoftRobot, -) +from soromox.systems.soft_robot import SoftRobot from soromox.systems.system_state import EnvironmentState, SystemState from .articulated import ArticulatedSoftRobot, McKibbenActuatedUMArm from .gvs import ( GVS, GVSSegment, - JointSpec, - LinkSpec, StrainBasisSpec, ) from .hsa import PlanarHSA @@ -65,6 +71,14 @@ "EnvironmentState", "SoftRobot", "CrossSectionGeometry", + "CrossSectionParams", + "ContinuumLinkParams", + "JointParams", + "IsotropicMaterialParams", + "LinearProfile", + "LinkSpec", + "JointSpec", + "shear_modulus_from_poisson_ratio", "SystemState", "custom_jvp_enabled", "custom_jvp_mode", @@ -81,15 +95,12 @@ # gvs systems "GVS", "GVSParams", - "GVSLinkParams", "GVSStructure", "GVSSegmentStructure", "GVSLinkStructure", "GVSJointStructure", "GVSStrainBasisStructure", "GVSSegment", - "LinkSpec", - "JointSpec", "StrainBasisSpec", # hsa systems "PlanarHSA", diff --git a/src/soromox/systems/articulated/articulated_soft_robot.py b/src/soromox/systems/articulated/articulated_soft_robot.py index adcaf14e..1ec43725 100644 --- a/src/soromox/systems/articulated/articulated_soft_robot.py +++ b/src/soromox/systems/articulated/articulated_soft_robot.py @@ -9,7 +9,8 @@ from soromox.actuation.core import Actuator, PassiveElement from soromox.systems.articulated.params import ArticulatedSoftRobotParams -from soromox.systems.soft_robot import CrossSectionGeometry, SoftRobot +from soromox.systems.components import CrossSectionGeometry +from soromox.systems.soft_robot import SoftRobot from soromox.utils.lie_algebra import se3, so3 diff --git a/src/soromox/systems/components/__init__.py b/src/soromox/systems/components/__init__.py new file mode 100644 index 00000000..53e8919d --- /dev/null +++ b/src/soromox/systems/components/__init__.py @@ -0,0 +1,31 @@ +"""Reusable link, joint, cross-section, and material components.""" + +from .cross_sections import ( + CrossSectionGeometry, + CrossSectionParams, + LinearProfile, + ProfileType, + evaluate_profile, + section_properties, +) +from .joints import JOINT_DOF, JointParams, JointSpec, JointType, joint_dof +from .links import ContinuumLinkParams, LinkSpec +from .materials import IsotropicMaterialParams, shear_modulus_from_poisson_ratio + +__all__ = [ + "ContinuumLinkParams", + "CrossSectionGeometry", + "CrossSectionParams", + "IsotropicMaterialParams", + "JOINT_DOF", + "JointParams", + "JointSpec", + "JointType", + "LinearProfile", + "LinkSpec", + "ProfileType", + "evaluate_profile", + "joint_dof", + "section_properties", + "shear_modulus_from_poisson_ratio", +] diff --git a/src/soromox/systems/components/cross_sections.py b/src/soromox/systems/components/cross_sections.py new file mode 100644 index 00000000..9fb26bfd --- /dev/null +++ b/src/soromox/systems/components/cross_sections.py @@ -0,0 +1,159 @@ +"""Shared cross-section parameterization and geometric properties.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import IntEnum +from typing import Literal + +import jax.numpy as jnp +from jax import Array + +from soromox.systems.params import BaseSystemParams + +__all__ = [ + "CrossSectionGeometry", + "CrossSectionParams", + "LinearProfile", + "ProfileType", + "evaluate_profile", + "section_properties", +] + + +class CrossSectionGeometry(IntEnum): + """Supported solid cross-section families.""" + + CIRCULAR = 0 + RECTANGULAR = 1 + ELLIPTICAL = 2 + + +ProfileType = Literal["constant", "linear"] + + +@dataclass(frozen=True) +class LinearProfile: + """A scalar cross-section dimension varying linearly along a link. + + Attributes: + base: Dimension at normalized arc length zero. + tip: Dimension at normalized arc length one. + """ + + base: float + tip: float + + +class CrossSectionParams(BaseSystemParams): + """Batched dynamic cross-section coefficients. + + Static geometry, profile kinds, and coefficient names live in each system's + structure. Rows are zero-padded to the largest coefficient count in the + model. + + Attributes: + coefficients: Cross-section coefficients with shape + ``(num_links, max_num_coefficients)``. Each row contains the + coefficients required by that link's static geometry and profile + declarations, followed by zero padding when necessary. + """ + + coefficients: Array + + def __check_init__(self) -> None: + self.validate() + + def validate(self) -> None: + """Validate the batched cross-section coefficient array. + + Returns: + None. + + Raises: + ValueError: If ``coefficients`` is not a nonempty two-dimensional + array. + """ + coefficients = jnp.asarray(self.coefficients) + if coefficients.ndim != 2: + raise ValueError( + "cross-section coefficients must have shape " + "(num_links, max_num_coefficients)." + ) + if coefficients.shape[0] < 1 or coefficients.shape[1] < 1: + raise ValueError("cross-section coefficients must be non-empty.") + + +def evaluate_profile( + coefficients: Array, + profile_type: ProfileType, + normalized_arclength: Array, +) -> Array: + """Evaluate a cross-section profile at normalized arc-length positions. + + Args: + coefficients: Coefficients for one profile. A constant profile uses + the first entry; a linear profile uses ``[base, tip]``. + profile_type: Either ``"constant"`` or ``"linear"``. + normalized_arclength: Scalar or array of positions measured from zero + at the link base to one at the link tip. + + Returns: + Profile values with the same shape as ``normalized_arclength``. + + Raises: + ValueError: If ``profile_type`` is not supported. + """ + coefficients = jnp.asarray(coefficients) + x = jnp.asarray(normalized_arclength) + if profile_type == "constant": + return jnp.broadcast_to(coefficients[0], x.shape) + if profile_type == "linear": + return coefficients[0] + x * (coefficients[1] - coefficients[0]) + raise ValueError(f"Unknown cross-section profile type {profile_type!r}.") + + +def section_properties( + geometry: CrossSectionGeometry | int, + dimensions: Array, +) -> tuple[Array, Array, Array, Array]: + """Return ``(I_x, I_y, I_z, area)`` for a solid section. + + The material-frame x-axis is longitudinal. Rectangular dimensions are + ordered ``[height, width]`` and elliptical dimensions are ordered + ``[semi_major, semi_minor]``. + + Args: + geometry: Cross-section family, supplied as a + :class:`CrossSectionGeometry` value or its integer value. + dimensions: Geometry dimensions. Circular sections use ``[radius]``; + rectangular sections use ``[height, width]``; elliptical sections + use ``[semi_major, semi_minor]``. + + Returns: + A tuple ``(I_x, I_y, I_z, area)`` containing the polar second moment + about the longitudinal material x-axis, the transverse second moments, + and the section area. + + Raises: + ValueError: If ``geometry`` does not identify a supported section. + """ + dimensions = jnp.asarray(dimensions) + geometry = CrossSectionGeometry(int(geometry)) + if geometry == CrossSectionGeometry.CIRCULAR: + radius = dimensions[0] + area = jnp.pi * radius**2 + transverse = jnp.pi * radius**4 / 4.0 + polar = 2.0 * transverse + return polar, transverse, transverse, area + if geometry == CrossSectionGeometry.RECTANGULAR: + height, width = dimensions[0], dimensions[1] + area = height * width + i_y = height * width**3 / 12.0 + i_z = width * height**3 / 12.0 + return i_y + i_z, i_y, i_z, area + semi_major, semi_minor = dimensions[0], dimensions[1] + area = jnp.pi * semi_major * semi_minor + i_y = jnp.pi * semi_major * semi_minor**3 / 4.0 + i_z = jnp.pi * semi_major**3 * semi_minor / 4.0 + return i_y + i_z, i_y, i_z, area diff --git a/src/soromox/systems/components/joints.py b/src/soromox/systems/components/joints.py new file mode 100644 index 00000000..c28de4c6 --- /dev/null +++ b/src/soromox/systems/components/joints.py @@ -0,0 +1,335 @@ +"""Shared joint construction specifications and runtime parameters.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Literal + +import jax.numpy as jnp +from jax import Array + +from soromox.systems.params import BaseSystemParams + +from .links import _validate_symmetric + +__all__ = ["JointParams", "JointSpec", "JointType", "JOINT_DOF", "joint_dof"] + +JointType = Literal[ + "revolute", + "prismatic", + "helical", + "cylindrical", + "planar", + "spherical", + "free", + "fixed", +] + +JOINT_DOF: dict[str, int] = { + "revolute": 1, + "prismatic": 1, + "helical": 1, + "cylindrical": 2, + "planar": 3, + "spherical": 3, + "free": 6, + "fixed": 0, +} + + +def joint_dof(joint_type: str) -> int: + """Return the active-coordinate dimension of a supported joint type. + + Args: + joint_type: Lower-case joint family name. + + Returns: + Number of active joint coordinates. + + Raises: + ValueError: If ``joint_type`` is unknown. + """ + try: + return JOINT_DOF[joint_type] + except KeyError as exc: + raise ValueError(f"Unknown joint type {joint_type!r}.") from exc + + +class JointParams(BaseSystemParams): + """Canonical batched generalized joint matrices. + + Attributes: + stiffness: Padded joint stiffness matrices with shape + ``(num_joints, generalized_dimension, generalized_dimension)``. + damping: Padded joint damping matrices with the same shape as + ``stiffness``. + """ + + stiffness: Array + damping: Array + + def __check_init__(self) -> None: + self.validate() + + def validate(self) -> None: + """Validate joint matrix shapes, finiteness, and symmetry. + + Returns: + None. + + Raises: + ValueError: If stiffness is not a batch of square matrices, damping + has a different shape, or either matrix batch is non-finite or + asymmetric. + """ + stiffness = jnp.asarray(self.stiffness) + damping = jnp.asarray(self.damping) + if stiffness.ndim != 3 or stiffness.shape[1] != stiffness.shape[2]: + raise ValueError( + "joint stiffness must have shape " + "(num_joints, generalized_dimension, generalized_dimension)." + ) + if damping.shape != stiffness.shape: + raise ValueError( + f"joint damping must have shape {stiffness.shape}, got {damping.shape}." + ) + _validate_symmetric("joint stiffness", stiffness) + _validate_symmetric("joint damping", damping) + + +@dataclass +class JointSpec: + """Kinematic and generalized viscoelastic specification of one joint. + + Attributes: + type: Joint family name. + axis: Axis used by revolute, prismatic, helical, and cylindrical joints. + plane: Motion plane used by planar joints. + pitch: Translation per radian for a helical joint. + stiffness: Stiffness matrix in active joint coordinates. An empty value + requests a zero matrix. + damping: Damping matrix in active joint coordinates. An empty value + requests a zero matrix. + """ + + type: JointType + axis: Literal["x", "y", "z"] = "x" + plane: Literal["xy", "yz", "xz"] = "xy" + pitch: float = 0.0 + stiffness: Array | list = field(default_factory=list) + damping: Array | list = field(default_factory=list) + + def __post_init__(self) -> None: + self.type = self.type.lower() # type: ignore[assignment] + dof = joint_dof(self.type) + if self.axis not in ("x", "y", "z"): + raise ValueError("axis must be one of 'x', 'y', or 'z'.") + if self.plane not in ("xy", "yz", "xz"): + raise ValueError("plane must be one of 'xy', 'yz', or 'xz'.") + for name in ("stiffness", "damping"): + value = jnp.asarray(getattr(self, name)) + if value.size == 0: + continue + if value.shape != (dof, dof): + raise ValueError( + f"joint {name} must have shape ({dof}, {dof}), got {value.shape}." + ) + _validate_symmetric(f"joint {name}", value) + + @classmethod + def fixed(cls) -> JointSpec: + """Create a zero-degree-of-freedom fixed joint. + + Returns: + A fixed-joint specification with zero stiffness and damping. + """ + return cls(type="fixed") + + @classmethod + def revolute( + cls, + axis: Literal["x", "y", "z"] = "x", + stiffness: Array | list | None = None, + damping: Array | list | None = None, + ) -> JointSpec: + """Create a one-degree-of-freedom revolute joint. + + Args: + axis: Rotation axis in the joint material frame. + stiffness: Optional ``(1, 1)`` stiffness matrix in active + coordinates. ``None`` produces zero stiffness. + damping: Optional ``(1, 1)`` damping matrix in active coordinates. + ``None`` produces zero damping. + + Returns: + A revolute-joint specification. + + Raises: + ValueError: If the axis is invalid or a matrix has an invalid shape, + non-finite entries, or asymmetry. + """ + return cls( + type="revolute", + axis=axis, + stiffness=[] if stiffness is None else stiffness, + damping=[] if damping is None else damping, + ) + + @classmethod + def prismatic( + cls, + axis: Literal["x", "y", "z"] = "x", + stiffness: Array | list | None = None, + damping: Array | list | None = None, + ) -> JointSpec: + """Create a one-degree-of-freedom prismatic joint. + + Args: + axis: Translation axis in the joint material frame. + stiffness: Optional ``(1, 1)`` active-coordinate stiffness matrix. + damping: Optional ``(1, 1)`` active-coordinate damping matrix. + + Returns: + A prismatic-joint specification. + + Raises: + ValueError: If the axis or either matrix is invalid. + """ + return cls( + type="prismatic", + axis=axis, + stiffness=[] if stiffness is None else stiffness, + damping=[] if damping is None else damping, + ) + + @classmethod + def helical( + cls, + axis: Literal["x", "y", "z"] = "x", + pitch: float = 0.0, + stiffness: Array | list | None = None, + damping: Array | list | None = None, + ) -> JointSpec: + """Create a one-degree-of-freedom helical joint. + + Args: + axis: Coupled rotation and translation axis. + pitch: Translation per radian of rotation. + stiffness: Optional ``(1, 1)`` active-coordinate stiffness matrix. + damping: Optional ``(1, 1)`` active-coordinate damping matrix. + + Returns: + A helical-joint specification. + + Raises: + ValueError: If the axis or either matrix is invalid. + """ + return cls( + type="helical", + axis=axis, + pitch=pitch, + stiffness=[] if stiffness is None else stiffness, + damping=[] if damping is None else damping, + ) + + @classmethod + def cylindrical( + cls, + axis: Literal["x", "y", "z"] = "x", + stiffness: Array | list | None = None, + damping: Array | list | None = None, + ) -> JointSpec: + """Create a cylindrical joint with rotation and translation. + + Args: + axis: Shared rotation and translation axis. + stiffness: Optional ``(2, 2)`` active-coordinate stiffness matrix. + damping: Optional ``(2, 2)`` active-coordinate damping matrix. + + Returns: + A two-degree-of-freedom cylindrical-joint specification. + + Raises: + ValueError: If the axis or either matrix is invalid. + """ + return cls( + type="cylindrical", + axis=axis, + stiffness=[] if stiffness is None else stiffness, + damping=[] if damping is None else damping, + ) + + @classmethod + def planar( + cls, + plane: Literal["xy", "yz", "xz"] = "xy", + stiffness: Array | list | None = None, + damping: Array | list | None = None, + ) -> JointSpec: + """Create a three-degree-of-freedom planar joint. + + Args: + plane: Plane in which the joint translates and rotates. + stiffness: Optional ``(3, 3)`` active-coordinate stiffness matrix. + damping: Optional ``(3, 3)`` active-coordinate damping matrix. + + Returns: + A planar-joint specification. + + Raises: + ValueError: If the plane or either matrix is invalid. + """ + return cls( + type="planar", + plane=plane, + stiffness=[] if stiffness is None else stiffness, + damping=[] if damping is None else damping, + ) + + @classmethod + def spherical( + cls, + stiffness: Array | list | None = None, + damping: Array | list | None = None, + ) -> JointSpec: + """Create a three-degree-of-freedom spherical joint. + + Args: + stiffness: Optional ``(3, 3)`` active-coordinate stiffness matrix. + damping: Optional ``(3, 3)`` active-coordinate damping matrix. + + Returns: + A spherical-joint specification. + + Raises: + ValueError: If either matrix is invalid. + """ + return cls( + type="spherical", + stiffness=[] if stiffness is None else stiffness, + damping=[] if damping is None else damping, + ) + + @classmethod + def free( + cls, + stiffness: Array | list | None = None, + damping: Array | list | None = None, + ) -> JointSpec: + """Create a free six-degree-of-freedom spatial joint. + + Args: + stiffness: Optional ``(6, 6)`` active-coordinate stiffness matrix. + damping: Optional ``(6, 6)`` active-coordinate damping matrix. + + Returns: + A free-joint specification. + + Raises: + ValueError: If either matrix is invalid. + """ + return cls( + type="free", + stiffness=[] if stiffness is None else stiffness, + damping=[] if damping is None else damping, + ) diff --git a/src/soromox/systems/components/links.py b/src/soromox/systems/components/links.py new file mode 100644 index 00000000..dfff03a1 --- /dev/null +++ b/src/soromox/systems/components/links.py @@ -0,0 +1,385 @@ +"""Shared continuum-link runtime parameters and construction specifications.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import jax.numpy as jnp +from jax import Array +from jax.errors import ConcretizationTypeError, TracerBoolConversionError + +from soromox.systems.params import BaseSystemParams + +from .cross_sections import ( + CrossSectionGeometry, + CrossSectionParams, + LinearProfile, + ProfileType, +) + +__all__ = ["ContinuumLinkParams", "LinkSpec"] + + +def _validate_symmetric(name: str, value: Array) -> None: + try: + finite = bool(jnp.all(jnp.isfinite(value))) + symmetric = bool(jnp.allclose(value, jnp.swapaxes(value, -1, -2))) + except (ConcretizationTypeError, TracerBoolConversionError): + return + if not finite: + raise ValueError(f"{name} must contain only finite values.") + if not symmetric: + raise ValueError(f"{name} must be symmetric in its trailing dimensions.") + + +class ContinuumLinkParams(BaseSystemParams): + """Canonical batched dynamic parameters for continuum links. + + Attributes: + length: Link lengths with shape ``(num_links,)``. + density: Volumetric mass densities with shape ``(num_links,)``. + reference_strain: Reference strain rows with shape + ``(num_links, strain_dimension)``. + cross_section: Batched cross-section coefficients shared by all + continuum-system families. + stiffness: Canonical generalized link stiffness matrices with shape + ``(num_links, generalized_dimension, generalized_dimension)``. + damping: Canonical generalized link damping matrices with the same + shape as ``stiffness``. + """ + + length: Array + density: Array + reference_strain: Array + cross_section: CrossSectionParams + stiffness: Array + damping: Array + + def __check_init__(self) -> None: + self.validate() + + def validate(self) -> None: + """Validate link-array shapes and canonical matrix properties. + + Returns: + None. + + Raises: + ValueError: If link fields disagree on ``num_links``, reference + strain or coefficient arrays are not two-dimensional, canonical + matrices are not equally shaped square batches, or a canonical + matrix contains non-finite values or is not symmetric. + """ + length = jnp.asarray(self.length) + if length.ndim != 1 or length.shape[0] < 1: + raise ValueError("length must have shape (num_links,) with num_links >= 1.") + num_links = length.shape[0] + density = jnp.asarray(self.density) + if density.shape != (num_links,): + raise ValueError( + f"density must have shape ({num_links},), got {density.shape}." + ) + reference_strain = jnp.asarray(self.reference_strain) + if reference_strain.ndim != 2 or reference_strain.shape[0] != num_links: + raise ValueError( + "reference_strain must have shape (num_links, strain_dimension)." + ) + self.cross_section.validate() + if self.cross_section.coefficients.shape[0] != num_links: + raise ValueError("cross-section coefficients must have one row per link.") + stiffness = jnp.asarray(self.stiffness) + damping = jnp.asarray(self.damping) + if stiffness.ndim != 3 or stiffness.shape[0] != num_links: + raise ValueError( + "stiffness must have shape " + "(num_links, generalized_dimension, generalized_dimension)." + ) + if stiffness.shape[1] != stiffness.shape[2]: + raise ValueError("stiffness trailing dimensions must be square.") + if damping.shape != stiffness.shape: + raise ValueError( + f"damping must have shape {stiffness.shape}, got {damping.shape}." + ) + _validate_symmetric("stiffness", stiffness) + _validate_symmetric("damping", damping) + + +def _profile( + value: float | LinearProfile, name: str +) -> tuple[tuple[float, ...], ProfileType, tuple[str, ...]]: + if isinstance(value, LinearProfile): + return (value.base, value.tip), "linear", (f"{name}_base", f"{name}_tip") + return (float(value),), "constant", (name,) + + +@dataclass(frozen=True) +class LinkSpec: + """Construction specification for one continuum link. + + Prefer the geometry-specific :meth:`circular`, :meth:`rectangular`, and + :meth:`elliptical` factories. Stiffness must be described by exactly one of + ``young_modulus`` plus ``shear_modulus`` or an explicit ``stiffness`` + matrix. Damping must similarly use exactly one of + ``material_damping_coefficient`` or an explicit ``damping`` matrix. + + Attributes: + cross_section_geometry: Static cross-section family. + length: Link length. + density: Volumetric mass density. + reference_strain: Reference strain vector for this link. + cross_section_coefficients: Packed constant or linear profile + coefficients. + cross_section_profile_types: Profile kind for each geometric dimension. + cross_section_profile_parameter_counts: Number of packed coefficients + occupied by each geometric dimension. + cross_section_coefficient_names: Human-readable names corresponding to + the packed coefficients. + young_modulus: Optional isotropic Young's modulus. + shear_modulus: Optional isotropic shear modulus. + material_damping_coefficient: Optional isotropic material damping value. + stiffness: Optional explicit generalized stiffness matrix. + damping: Optional explicit generalized damping matrix. + """ + + cross_section_geometry: CrossSectionGeometry + length: float + density: float + reference_strain: Array | list[float] + cross_section_coefficients: tuple[float, ...] + cross_section_profile_types: tuple[ProfileType, ...] + cross_section_profile_parameter_counts: tuple[int, ...] + cross_section_coefficient_names: tuple[str, ...] + young_modulus: float | None = None + shear_modulus: float | None = None + material_damping_coefficient: float | None = None + stiffness: Array | None = None + damping: Array | None = None + + def __post_init__(self) -> None: + material_stiffness = ( + self.young_modulus is not None or self.shear_modulus is not None + ) + if material_stiffness and self.stiffness is not None: + raise ValueError( + "Provide either young_modulus/shear_modulus or stiffness, not both." + ) + if not material_stiffness and self.stiffness is None: + raise ValueError( + "Provide young_modulus and shear_modulus or an explicit stiffness." + ) + if material_stiffness and ( + self.young_modulus is None or self.shear_modulus is None + ): + raise ValueError( + "Material stiffness construction requires both young_modulus " + "and shear_modulus." + ) + if self.material_damping_coefficient is not None and self.damping is not None: + raise ValueError( + "Provide either material_damping_coefficient or damping, not both." + ) + if self.material_damping_coefficient is None and self.damping is None: + raise ValueError( + "Provide material_damping_coefficient or an explicit damping matrix." + ) + for name in ("stiffness", "damping"): + value = getattr(self, name) + if value is None: + continue + matrix = jnp.asarray(value) + if matrix.ndim != 2 or matrix.shape[0] != matrix.shape[1]: + raise ValueError(f"Explicit {name} must be a square matrix.") + _validate_symmetric(name, matrix) + + @classmethod + def _make( + cls, + *, + geometry: CrossSectionGeometry, + profiles: tuple[tuple[str, float | LinearProfile], ...], + length: float, + density: float, + reference_strain: Array | list[float], + young_modulus: float | None, + shear_modulus: float | None, + material_damping_coefficient: float | None, + stiffness: Array | None, + damping: Array | None, + ) -> LinkSpec: + packed = [_profile(value, name) for name, value in profiles] + return cls( + cross_section_geometry=geometry, + length=length, + density=density, + reference_strain=reference_strain, + cross_section_coefficients=tuple( + coefficient for values, _, _ in packed for coefficient in values + ), + cross_section_profile_types=tuple(kind for _, kind, _ in packed), + cross_section_profile_parameter_counts=tuple( + len(values) for values, _, _ in packed + ), + cross_section_coefficient_names=tuple( + name for _, _, names in packed for name in names + ), + young_modulus=young_modulus, + shear_modulus=shear_modulus, + material_damping_coefficient=material_damping_coefficient, + stiffness=stiffness, + damping=damping, + ) + + @classmethod + def circular( + cls, + *, + length: float, + radius: float | LinearProfile, + density: float, + reference_strain: Array | list[float], + young_modulus: float | None = None, + shear_modulus: float | None = None, + material_damping_coefficient: float | None = None, + stiffness: Array | None = None, + damping: Array | None = None, + ) -> LinkSpec: + """Create a solid circular-link specification. + + Args: + length: Link length. + radius: Constant radius or a base-to-tip linear radius profile. + density: Volumetric mass density. + reference_strain: Reference strain vector. Spatial systems use six + entries and planar PCS uses three entries. + young_modulus: Young's modulus used with ``shear_modulus`` to build + generalized stiffness. Mutually exclusive with ``stiffness``. + shear_modulus: Shear modulus used with ``young_modulus``. + material_damping_coefficient: Isotropic damping value used to build + generalized damping. Mutually exclusive with ``damping``. + stiffness: Explicit generalized stiffness matrix. + damping: Explicit generalized damping matrix. + + Returns: + A circular :class:`LinkSpec` ready for PCS or GVS construction. + + Raises: + ValueError: If material and explicit sources are missing, + incomplete, or supplied together, or if an explicit matrix is + not finite, square, and symmetric. + """ + return cls._make( + geometry=CrossSectionGeometry.CIRCULAR, + profiles=(("radius", radius),), + length=length, + density=density, + reference_strain=reference_strain, + young_modulus=young_modulus, + shear_modulus=shear_modulus, + material_damping_coefficient=material_damping_coefficient, + stiffness=stiffness, + damping=damping, + ) + + @classmethod + def rectangular( + cls, + *, + length: float, + height: float | LinearProfile, + width: float | LinearProfile, + density: float, + reference_strain: Array | list[float], + young_modulus: float | None = None, + shear_modulus: float | None = None, + material_damping_coefficient: float | None = None, + stiffness: Array | None = None, + damping: Array | None = None, + ) -> LinkSpec: + """Create a solid rectangular-link specification. + + Args: + length: Link length. + height: Constant height or base-to-tip linear height profile. + width: Constant width or base-to-tip linear width profile. + density: Volumetric mass density. + reference_strain: Reference strain vector for the link. + young_modulus: Young's modulus used with ``shear_modulus`` to build + generalized stiffness. Mutually exclusive with ``stiffness``. + shear_modulus: Shear modulus used with ``young_modulus``. + material_damping_coefficient: Isotropic damping value used to build + generalized damping. Mutually exclusive with ``damping``. + stiffness: Explicit generalized stiffness matrix. + damping: Explicit generalized damping matrix. + + Returns: + A rectangular :class:`LinkSpec` ready for GVS construction. + + Raises: + ValueError: If material and explicit sources are missing, + incomplete, or supplied together, or if an explicit matrix is + not finite, square, and symmetric. + """ + return cls._make( + geometry=CrossSectionGeometry.RECTANGULAR, + profiles=(("height", height), ("width", width)), + length=length, + density=density, + reference_strain=reference_strain, + young_modulus=young_modulus, + shear_modulus=shear_modulus, + material_damping_coefficient=material_damping_coefficient, + stiffness=stiffness, + damping=damping, + ) + + @classmethod + def elliptical( + cls, + *, + length: float, + semi_major: float | LinearProfile, + semi_minor: float | LinearProfile, + density: float, + reference_strain: Array | list[float], + young_modulus: float | None = None, + shear_modulus: float | None = None, + material_damping_coefficient: float | None = None, + stiffness: Array | None = None, + damping: Array | None = None, + ) -> LinkSpec: + """Create a solid elliptical-link specification. + + Args: + length: Link length. + semi_major: Constant semi-major axis or base-to-tip linear profile. + semi_minor: Constant semi-minor axis or base-to-tip linear profile. + density: Volumetric mass density. + reference_strain: Reference strain vector for the link. + young_modulus: Young's modulus used with ``shear_modulus`` to build + generalized stiffness. Mutually exclusive with ``stiffness``. + shear_modulus: Shear modulus used with ``young_modulus``. + material_damping_coefficient: Isotropic damping value used to build + generalized damping. Mutually exclusive with ``damping``. + stiffness: Explicit generalized stiffness matrix. + damping: Explicit generalized damping matrix. + + Returns: + An elliptical :class:`LinkSpec` ready for GVS construction. + + Raises: + ValueError: If material and explicit sources are missing, + incomplete, or supplied together, or if an explicit matrix is + not finite, square, and symmetric. + """ + return cls._make( + geometry=CrossSectionGeometry.ELLIPTICAL, + profiles=(("semi_major", semi_major), ("semi_minor", semi_minor)), + length=length, + density=density, + reference_strain=reference_strain, + young_modulus=young_modulus, + shear_modulus=shear_modulus, + material_damping_coefficient=material_damping_coefficient, + stiffness=stiffness, + damping=damping, + ) diff --git a/src/soromox/systems/components/materials.py b/src/soromox/systems/components/materials.py new file mode 100644 index 00000000..3bd70e40 --- /dev/null +++ b/src/soromox/systems/components/materials.py @@ -0,0 +1,103 @@ +"""Shared isotropic material construction and optimization parameters.""" + +from __future__ import annotations + +import jax.numpy as jnp +from jax import Array + +from soromox.systems.params import BaseSystemParams + +__all__ = ["IsotropicMaterialParams", "shear_modulus_from_poisson_ratio"] + + +def shear_modulus_from_poisson_ratio( + young_modulus: Array | float, + poisson_ratio: Array | float, +) -> Array: + """Compute the shear modulus of an isotropic material. + + Args: + young_modulus: Young's modulus ``E`` as a scalar or array. + poisson_ratio: Poisson's ratio ``nu`` as a scalar or array broadcastable + with ``young_modulus``. + + Returns: + The shear modulus ``G = E / (2 (1 + nu))`` with the broadcasted input + shape. + """ + young = jnp.asarray(young_modulus) + poisson = jnp.asarray(poisson_ratio) + return young / (2.0 * (1.0 + poisson)) + + +class IsotropicMaterialParams(BaseSystemParams): + """Caller-owned isotropic material variables. + + These values are construction and optimization variables rather than + canonical robot runtime parameters. Pass an instance to + ``link_matrices_from_material`` or ``with_isotropic_material`` to map it to + generalized link matrices. + + Attributes: + young_modulus: Scalar or per-link Young's modulus values. + shear_modulus: Scalar or per-link shear modulus values. + material_damping_coefficient: Scalar or per-link isotropic material + damping values. + """ + + young_modulus: Array + shear_modulus: Array + material_damping_coefficient: Array + + def __check_init__(self) -> None: + self.validate() + + def validate(self) -> None: + """Validate material field dimensionality. + + Returns: + None. + + Raises: + ValueError: If any material field has more than one dimension. + """ + for name in ( + "young_modulus", + "shear_modulus", + "material_damping_coefficient", + ): + value = jnp.asarray(getattr(self, name)) + if value.ndim > 1: + raise ValueError(f"{name} must be scalar or one-dimensional.") + + def broadcast(self, num_links: int) -> IsotropicMaterialParams: + """Broadcast scalar material fields to one value per link. + + Args: + num_links: Number of links in the target robot. + + Returns: + A new material PyTree whose three fields all have shape + ``(num_links,)``. + + Raises: + ValueError: If a non-scalar field does not have shape + ``(num_links,)``. + """ + + def broadcast_one(name: str) -> Array: + value = jnp.asarray(getattr(self, name)) + if value.ndim == 0: + return jnp.full((num_links,), value) + if value.shape != (num_links,): + raise ValueError( + f"{name} must be scalar or have shape ({num_links},), " + f"got {value.shape}." + ) + return value + + return IsotropicMaterialParams( + young_modulus=broadcast_one("young_modulus"), + shear_modulus=broadcast_one("shear_modulus"), + material_damping_coefficient=broadcast_one("material_damping_coefficient"), + ) diff --git a/src/soromox/systems/gvs/__init__.py b/src/soromox/systems/gvs/__init__.py index 48d06c0b..85c22751 100644 --- a/src/soromox/systems/gvs/__init__.py +++ b/src/soromox/systems/gvs/__init__.py @@ -1,7 +1,4 @@ -from soromox.systems.gvs.params import ( - GVSLinkParams, - GVSParams, -) +from soromox.systems.gvs.params import GVSParams from soromox.systems.gvs.structures import ( GVSJointStructure, GVSLinkStructure, @@ -11,15 +8,12 @@ ) from .core import GVS -from .specs import GVSSegment, JointSpec, LinkSpec, StrainBasisSpec +from .specs import GVSSegment, StrainBasisSpec __all__ = [ "GVS", "GVSSegment", - "LinkSpec", - "JointSpec", "StrainBasisSpec", - "GVSLinkParams", "GVSParams", "GVSStructure", "GVSSegmentStructure", diff --git a/src/soromox/systems/gvs/_assembly.py b/src/soromox/systems/gvs/_assembly.py index 1b8678b4..a497179c 100644 --- a/src/soromox/systems/gvs/_assembly.py +++ b/src/soromox/systems/gvs/_assembly.py @@ -16,6 +16,24 @@ def _set_model_field(model: Any, name: str, value: Any) -> None: object.__setattr__(model, name, value) +def _profile_endpoint_pairs(segment, coefficients): + pairs = [] + start = 0 + for profile_type, count in zip( + segment.link.cross_section_profile_types, + segment.link.cross_section_profile_parameter_counts, + strict=True, + ): + values = coefficients[start : start + count] + pairs.append( + (values[0], values[0]) + if profile_type == "constant" + else (values[0], values[1]) + ) + start += count + return pairs + + def _validate_segments( segments: list[GVSSegmentStructure] | tuple[GVSSegmentStructure, ...], params: GVSParams, @@ -72,7 +90,10 @@ def assign_gvs_runtime_arrays( dofs_link = [ int( Basis.DOF_BRANCHES[Basis.BASISTYPE_MAP[s.basis.type]]( - (jnp.asarray(s.basis.active), jnp.asarray(s.basis.orders)) + ( + jnp.asarray(s.basis.strain_selector), + jnp.asarray(s.basis.basis_order), + ) ) ) for s in segments @@ -104,13 +125,6 @@ def assign_gvs_runtime_arrays( mass_matrices = jnp.empty( (n_segments, max_num_integration_points, 6, 6), dtype=float ) - stiffness_matrices = jnp.empty( - (n_segments, max_num_integration_points, 6, 6), dtype=float - ) - damping_matrices = jnp.empty( - (n_segments, max_num_integration_points, 6, 6), dtype=float - ) - B_joint = jnp.empty((n_segments, 6, max_dof), dtype=float) B_Xs = jnp.empty((n_segments, max_num_integration_points, 6, max_dof), dtype=float) B_Z1 = jnp.empty( @@ -139,13 +153,14 @@ def assign_gvs_runtime_arrays( for i_segment, segment in enumerate(segments): joint_dof = Joint.DICT_JOINT_TYPE_DOF[segment.joint.type] - if ( - params.joint_stiffness.shape[1] < joint_dof - or params.joint_stiffness.shape[2] < joint_dof - ): - raise ValueError( - "joint_stiffness trailing dimensions must cover every joint DOF." - ) + pairs = _profile_endpoint_pairs( + segment, params.link.cross_section.coefficients[i_segment] + ) + zero_pair = (jnp.array(0.0), jnp.array(0.0)) + geometry = int(segment.link.cross_section_geometry) + radius_pair = pairs[0] if geometry == 0 else zero_pair + rectangular_pairs = pairs if geometry == 1 else [zero_pair, zero_pair] + elliptical_pairs = pairs if geometry == 2 else [zero_pair, zero_pair] segment_data = model._build_segment_i( max_dof=max_dof, max_num_integration_points=max_num_integration_points, @@ -154,22 +169,19 @@ def assign_gvs_runtime_arrays( basis_structure=segment.basis, num_gauss_points=segment.num_gauss_points, length=params.link.length[i_segment], - young_modulus=params.link.young_modulus[i_segment], - poisson_ratio=params.link.poisson_ratio[i_segment], density=params.link.density[i_segment], - damping_coefficient=params.link.damping_coefficient[i_segment], - radius_initial=params.link.radius_initial[i_segment], - radius_final=params.link.radius_final[i_segment], - height_initial=params.link.height_initial[i_segment], - height_final=params.link.height_final[i_segment], - width_initial=params.link.width_initial[i_segment], - width_final=params.link.width_final[i_segment], - semi_major_initial=params.link.semi_major_initial[i_segment], - semi_major_final=params.link.semi_major_final[i_segment], - semi_minor_initial=params.link.semi_minor_initial[i_segment], - semi_minor_final=params.link.semi_minor_final[i_segment], - reference_strain=params.reference_strain[i_segment], - joint_stiffness=params.joint_stiffness[i_segment, :joint_dof, :joint_dof], + radius_initial=radius_pair[0], + radius_final=radius_pair[1], + height_initial=rectangular_pairs[0][0], + height_final=rectangular_pairs[0][1], + width_initial=rectangular_pairs[1][0], + width_final=rectangular_pairs[1][1], + semi_major_initial=elliptical_pairs[0][0], + semi_major_final=elliptical_pairs[0][1], + semi_minor_initial=elliptical_pairs[1][0], + semi_minor_final=elliptical_pairs[1][1], + reference_strain=params.link.reference_strain[i_segment], + joint_stiffness=params.joint.stiffness[i_segment, :joint_dof, :joint_dof], ) segment_lengths = segment_lengths.at[i_segment].set(segment_data.L) @@ -186,12 +198,6 @@ def assign_gvs_runtime_arrays( segment_data.integration_weights ) mass_matrices = mass_matrices.at[i_segment].set(segment_data.mass_matrices) - stiffness_matrices = stiffness_matrices.at[i_segment].set( - segment_data.stiffness_matrices - ) - damping_matrices = damping_matrices.at[i_segment].set( - segment_data.damping_matrices - ) B_joint = B_joint.at[i_segment].set(segment_data.B_joint) B_Xs = B_Xs.at[i_segment].set(segment_data.B_Xs) B_Z1 = B_Z1.at[i_segment].set(segment_data.B_Z1) @@ -211,32 +217,18 @@ def assign_gvs_runtime_arrays( basis_type_index = basis_type_index.at[i_segment].set( Basis.BASISTYPE_MAP[basis.type] ) - basis_active_params = basis_active_params.at[i_segment].set(basis.active) - basis_order_params = basis_order_params.at[i_segment].set(basis.orders) + basis_active_params = basis_active_params.at[i_segment].set( + basis.strain_selector + ) + basis_order_params = basis_order_params.at[i_segment].set(basis.basis_order) cross_section_geometry = cross_section_geometry.at[i_segment].set( int(segment.link.cross_section_geometry) ) - radius_params = radius_params.at[i_segment].set( - [params.link.radius_initial[i_segment], params.link.radius_final[i_segment]] - ) - height_params = height_params.at[i_segment].set( - [params.link.height_initial[i_segment], params.link.height_final[i_segment]] - ) - width_params = width_params.at[i_segment].set( - [params.link.width_initial[i_segment], params.link.width_final[i_segment]] - ) - semi_major_params = semi_major_params.at[i_segment].set( - [ - params.link.semi_major_initial[i_segment], - params.link.semi_major_final[i_segment], - ] - ) - semi_minor_params = semi_minor_params.at[i_segment].set( - [ - params.link.semi_minor_initial[i_segment], - params.link.semi_minor_final[i_segment], - ] - ) + radius_params = radius_params.at[i_segment].set(radius_pair) + height_params = height_params.at[i_segment].set(rectangular_pairs[0]) + width_params = width_params.at[i_segment].set(rectangular_pairs[1]) + semi_major_params = semi_major_params.at[i_segment].set(elliptical_pairs[0]) + semi_minor_params = semi_minor_params.at[i_segment].set(elliptical_pairs[1]) _set_model_field(model, "segment_lengths", segment_lengths) _set_model_field(model, "num_integration_points", num_integration_points) @@ -244,8 +236,6 @@ def assign_gvs_runtime_arrays( _set_model_field(model, "integration_points", integration_points) _set_model_field(model, "integration_weights", integration_weights) _set_model_field(model, "mass_matrices", mass_matrices) - _set_model_field(model, "stiffness_matrices", stiffness_matrices) - _set_model_field(model, "damping_matrices", damping_matrices) _set_model_field(model, "B_joint", B_joint) _set_model_field(model, "B_Xs", B_Xs) _set_model_field(model, "B_Z1", B_Z1) diff --git a/src/soromox/systems/gvs/_runtime.py b/src/soromox/systems/gvs/_runtime.py index 36b19863..dd57786b 100644 --- a/src/soromox/systems/gvs/_runtime.py +++ b/src/soromox/systems/gvs/_runtime.py @@ -1,4 +1,4 @@ -__all__ = ["LinkRuntimeData", "SegmentRuntimeData"] +__all__ = ["SegmentRuntimeData"] from dataclasses import dataclass @@ -18,8 +18,6 @@ class SegmentRuntimeData: integration_points: Array integration_weights: Array mass_matrices: Array - stiffness_matrices: Array - damping_matrices: Array B_joint: Array B_Xs: Array B_Z1: Array @@ -39,8 +37,6 @@ def tree_flatten(self): self.integration_points, self.integration_weights, self.mass_matrices, - self.stiffness_matrices, - self.damping_matrices, self.B_joint, self.B_Xs, self.B_Z1, @@ -56,49 +52,3 @@ def tree_flatten(self): @classmethod def tree_unflatten(cls, aux_data, children): return cls(*children) - - -@jax.tree_util.register_pytree_node_class -@dataclass -class LinkRuntimeData: - """Discretized per-link runtime data for link-only computations.""" - - L: Array - num_integration_points: Array - strain_selector: Array - integration_points: Array - integration_weights: Array - mass_matrices: Array - stiffness_matrices: Array - damping_matrices: Array - B_Xs: Array - B_Z1: Array - B_Z2: Array - xi_ref_Xs: Array - xi_ref_Z1: Array - xi_ref_Z2: Array - dof_link: Array - - def tree_flatten(self): - children = ( - self.L, - self.num_integration_points, - self.strain_selector, - self.integration_points, - self.integration_weights, - self.mass_matrices, - self.stiffness_matrices, - self.damping_matrices, - self.B_Xs, - self.B_Z1, - self.B_Z2, - self.xi_ref_Xs, - self.xi_ref_Z1, - self.xi_ref_Z2, - self.dof_link, - ) - return children, None - - @classmethod - def tree_unflatten(cls, aux_data, children): - return cls(*children) diff --git a/src/soromox/systems/gvs/construction.py b/src/soromox/systems/gvs/construction.py index 9c36c696..23fb4ea1 100644 --- a/src/soromox/systems/gvs/construction.py +++ b/src/soromox/systems/gvs/construction.py @@ -1,11 +1,28 @@ -__all__ = ["params_and_structure_from_segments"] +"""Construction and material projection for GVS segment specifications.""" -from jax import Array +from __future__ import annotations + +from jax import Array, vmap from jax import numpy as jnp -from soromox.systems.gvs.params import GVSLinkParams, GVSParams +from soromox.systems.components import ( + ContinuumLinkParams, + CrossSectionParams, + JointParams, + evaluate_profile, + section_properties, +) +from soromox.systems.gvs.params import GVSParams from soromox.systems.gvs.primitives import Basis, Joint from soromox.systems.gvs.specs import GVSSegment +from soromox.systems.gvs.strain_bases import ( + B_IMQ, + B_Chebychev, + B_Fourier, + B_Gaussian, + B_LegendrePolynomial, + B_Monomial, +) from soromox.systems.gvs.structures import ( GVSJointStructure, GVSLinkStructure, @@ -13,6 +30,153 @@ GVSStrainBasisStructure, GVSStructure, ) +from soromox.utils.integration import gauss_quadrature + +__all__ = ["material_operators_from_params", "params_and_structure_from_segments"] + +_BASIS_FUNCTIONS = { + "monomial": B_Monomial, + "legendre": B_LegendrePolynomial, + "chebyshev": B_Chebychev, + "fourier": B_Fourier, + "gaussian": B_Gaussian, + "imq": B_IMQ, +} + + +def _link_dof(segment: GVSSegmentStructure) -> int: + return int( + Basis.DOF_BRANCHES[Basis.BASISTYPE_MAP[segment.basis.type]]( + ( + jnp.asarray(segment.basis.strain_selector), + jnp.asarray(segment.basis.basis_order), + ) + ) + ) + + +def _basis_at_points( + segment: GVSSegmentStructure, points: Array, max_dof: int +) -> Array: + function = _BASIS_FUNCTIONS[segment.basis.type] + selector = jnp.asarray(segment.basis.strain_selector) + order = jnp.asarray(segment.basis.basis_order) + return vmap(lambda point: function(point, selector, order, max_dof))(points) + + +def _dimensions_at_points( + coefficients: Array, structure: GVSLinkStructure, points: Array +) -> Array: + dimensions = [] + start = 0 + for profile_type, count in zip( + structure.cross_section_profile_types, + structure.cross_section_profile_parameter_counts, + strict=True, + ): + stop = start + count + dimensions.append( + evaluate_profile(coefficients[start:stop], profile_type, points) + ) + start = stop + return jnp.stack(dimensions, axis=1) + + +def material_operators_from_params( + params: GVSParams, structure: GVSStructure +) -> tuple[Array, Array, Array]: + """Build unit-response material operators for all GVS links. + + The returned arrays project cross-section mechanics through each link's + configured strain basis and Gauss quadrature rule. They satisfy + ``K = E * K_E + G * K_G`` and ``D = eta * D_eta`` when the per-link + material scalars are broadcast over the final two axes. + + Args: + params: Canonical GVS parameters providing link lengths and packed + cross-section coefficients. + structure: Static GVS layout providing cross-section geometries, + profile metadata, bases, quadrature rules, and the padded DOF size. + + Returns: + A tuple ``(young_operators, shear_operators, damping_operators)``. Each + array has shape ``(num_segments, max_dof, max_dof)`` and represents the + response to a unit Young's modulus, shear modulus, or material damping + coefficient, respectively. + + Raises: + ValueError: If the parameters or structure contain incompatible profile + coefficients, unsupported cross-section data, or invalid quadrature + configuration. + """ + max_dof = int(structure.max_dof) + young_operators = [] + shear_operators = [] + damping_operators = [] + for index, segment in enumerate(structure.segments): + points, weights, _ = gauss_quadrature(segment.num_gauss_points) + basis = _basis_at_points(segment, points, max_dof) + length = params.link.length[index] + if structure.scale_rotational_basis_by_length: + basis = basis.at[:, :3, :].divide(length) + dimensions = _dimensions_at_points( + params.link.cross_section.coefficients[index], segment.link, points + ) + geometry = segment.link.cross_section_geometry + ix, iy, iz, area = vmap( + lambda value, geometry=geometry: section_properties(geometry, value) + )(dimensions) + young_diagonal = jnp.stack( + [jnp.zeros_like(ix), iy, iz, area, jnp.zeros_like(ix), jnp.zeros_like(ix)], + axis=1, + ) + shear_diagonal = jnp.stack( + [ + ix, + jnp.zeros_like(ix), + jnp.zeros_like(ix), + jnp.zeros_like(ix), + area, + area, + ], + axis=1, + ) + damping_diagonal = jnp.stack( + [ix, 3.0 * iy, 3.0 * iz, 3.0 * area, area, area], axis=1 + ) + + def project( + diagonal: Array, + basis: Array = basis, + length: Array = length, + weights: Array = weights, + ) -> Array: + local = vmap(jnp.diag)(diagonal) + integrand = vmap(lambda b, c: b.T @ c @ b)(basis, local) + return length * jnp.sum(weights[:, None, None] * integrand, axis=0) + + young_operators.append(project(young_diagonal)) + shear_operators.append(project(shear_diagonal)) + damping_operators.append(project(damping_diagonal)) + return ( + jnp.stack(young_operators), + jnp.stack(shear_operators), + jnp.stack(damping_operators), + ) + + +def _pad_matrix(value: Array | list, dof: int, max_dof: int, name: str) -> Array: + matrix = jnp.asarray(value) + if matrix.size == 0: + return jnp.zeros((max_dof, max_dof)) + if matrix.shape == (max_dof, max_dof): + return matrix + if matrix.shape != (dof, dof): + raise ValueError( + f"{name} must have shape ({dof}, {dof}) or " + f"({max_dof}, {max_dof}), got {matrix.shape}." + ) + return jnp.pad(matrix, ((0, max_dof - dof), (0, max_dof - dof))) def params_and_structure_from_segments( @@ -24,80 +188,49 @@ def params_and_structure_from_segments( max_num_gauss_points: int | None = None, scale_rotational_basis_by_length: bool = False, ) -> tuple[GVSParams, GVSStructure]: - """Split user-facing GVS segment specs into dynamic params and structure.""" + """Convert segment specifications to canonical params and static structure. + + Link and joint matrices expressed in active coordinates are zero-padded to + the common ``max_dof`` size. Isotropic link material properties are first + projected through the configured strain bases; explicit link matrices pass + through unchanged apart from padding. + + Args: + segments: Non-empty sequence of segment construction specifications. + gravity: Optional world-frame gravity vector with shape ``(3,)``. + Defaults to ``[0, 0, 9.81]``. + base_pose: Optional base translation and quaternion with shape ``(7,)``. + Defaults to the origin with an identity quaternion. + max_dof: Optional common padded generalized-coordinate dimension. When + omitted, the largest joint or link DOF among the segments is used. + max_num_gauss_points: Optional common quadrature padding size. When + omitted, the largest segment quadrature rule is used. + scale_rotational_basis_by_length: Whether rotational strain-basis rows + are divided by the associated link length during projection. + + Returns: + A tuple ``(params, structure)`` containing validated canonical + :class:`GVSParams` and the corresponding static :class:`GVSStructure`. + + Raises: + KeyError: If a joint or basis type is unsupported. + ValueError: If ``segments`` is empty; a quadrature rule has fewer than + five points; an explicit matrix has an invalid shape; a material + source is incomplete; or either padding size is too small. + """ if not segments: raise ValueError("GVS requires at least one segment.") input_segments = tuple(segments) - n_segments = len(input_segments) - dofs_joint = [ - Joint.DICT_JOINT_TYPE_DOF[segment.joint.type] for segment in input_segments - ] - dofs_link = [] - for segment in input_segments: - basis_type_idx = Basis.BASISTYPE_MAP[segment.basis.type] - dofs_link.append( - int( - Basis.DOF_BRANCHES[basis_type_idx]( - ( - jnp.asarray(segment.basis.active), - jnp.asarray(segment.basis.orders), - ) - ) - ) - ) - inferred_max_dof = max(dofs_joint + dofs_link) - layout_max_dof = inferred_max_dof if max_dof is None else max_dof - if layout_max_dof < inferred_max_dof: - raise ValueError( - f"max_dof={layout_max_dof} is smaller than the required " - f"GVS segment DOF {inferred_max_dof}." - ) + joint_dofs = [Joint.DICT_JOINT_TYPE_DOF[s.joint.type] for s in input_segments] - joint_stiffness = jnp.zeros((n_segments, layout_max_dof, layout_max_dof)) - for i, segment in enumerate(input_segments): - dof = dofs_joint[i] - stiffness = jnp.asarray(segment.joint.stiffness) - if dof == 0: - continue - if stiffness.size == 0: - continue - if stiffness.shape != (dof, dof): - raise ValueError( - "joint stiffness must have shape " - f"({dof}, {dof}) for segment {i}, got {stiffness.shape}." - ) - joint_stiffness = joint_stiffness.at[i, :dof, :dof].set(stiffness) - - links = [segment.link for segment in input_segments] - params = GVSParams( - link=GVSLinkParams( - length=jnp.asarray([link.L for link in links]), - young_modulus=jnp.asarray([link.E for link in links]), - poisson_ratio=jnp.asarray([link.nu for link in links]), - density=jnp.asarray([link.rho for link in links]), - damping_coefficient=jnp.asarray([link.eta for link in links]), - radius_initial=jnp.asarray([link.r_i for link in links]), - radius_final=jnp.asarray([link.r_f for link in links]), - height_initial=jnp.asarray([link.h_i for link in links]), - height_final=jnp.asarray([link.h_f for link in links]), - width_initial=jnp.asarray([link.w_i for link in links]), - width_final=jnp.asarray([link.w_f for link in links]), - semi_major_initial=jnp.asarray([link.a_i for link in links]), - semi_major_final=jnp.asarray([link.a_f for link in links]), - semi_minor_initial=jnp.asarray([link.b_i for link in links]), - semi_minor_final=jnp.asarray([link.b_f for link in links]), - ), - gravity=gravity, - base_pose=base_pose, - reference_strain=jnp.asarray( - [segment.basis.xi_ref for segment in input_segments] - ), - joint_stiffness=joint_stiffness, - ) static_segments = tuple( GVSSegmentStructure( link=GVSLinkStructure( - cross_section_geometry=segment.link.cross_section_geometry + cross_section_geometry=segment.link.cross_section_geometry, + cross_section_profile_types=segment.link.cross_section_profile_types, + cross_section_profile_parameter_counts=( + segment.link.cross_section_profile_parameter_counts + ), ), joint=GVSJointStructure( type=segment.joint.type, @@ -107,23 +240,114 @@ def params_and_structure_from_segments( ), basis=GVSStrainBasisStructure( type=segment.basis.type, - active=tuple( + strain_selector=tuple( int(value) - for value in jnp.asarray(segment.basis.active).reshape(-1).tolist() + for value in jnp.asarray(segment.basis.strain_selector) + .reshape(-1) + .tolist() ), - orders=tuple( + basis_order=tuple( int(value) - for value in jnp.asarray(segment.basis.orders).reshape(-1).tolist() + for value in jnp.asarray(segment.basis.basis_order) + .reshape(-1) + .tolist() ), ), num_gauss_points=int(segment.num_gauss_points), ) for segment in input_segments ) + link_dofs = [_link_dof(segment) for segment in static_segments] + required_max_dof = max(joint_dofs + link_dofs) + layout_max_dof = required_max_dof if max_dof is None else max_dof + if layout_max_dof < required_max_dof: + raise ValueError( + f"max_dof={layout_max_dof} is smaller than required DOF {required_max_dof}." + ) structure = GVSStructure( segments=static_segments, max_dof=layout_max_dof, max_num_gauss_points=max_num_gauss_points, scale_rotational_basis_by_length=scale_rotational_basis_by_length, ) + + max_coefficients = max( + len(segment.link.cross_section_coefficients) for segment in input_segments + ) + coefficients = jnp.zeros((len(input_segments), max_coefficients)) + for index, segment in enumerate(input_segments): + values = jnp.asarray(segment.link.cross_section_coefficients) + coefficients = coefficients.at[index, : values.size].set(values) + + zeros = jnp.zeros((len(input_segments), layout_max_dof, layout_max_dof)) + joint_stiffness = [] + joint_damping = [] + for index, segment in enumerate(input_segments): + joint_stiffness.append( + _pad_matrix( + segment.joint.stiffness, + joint_dofs[index], + layout_max_dof, + "joint stiffness", + ) + ) + joint_damping.append( + _pad_matrix( + segment.joint.damping, + joint_dofs[index], + layout_max_dof, + "joint damping", + ) + ) + params = GVSParams( + base_pose=base_pose, + gravity=gravity, + link=ContinuumLinkParams( + length=jnp.asarray([segment.link.length for segment in input_segments]), + density=jnp.asarray([segment.link.density for segment in input_segments]), + reference_strain=jnp.asarray( + [segment.link.reference_strain for segment in input_segments], + dtype=float, + ), + cross_section=CrossSectionParams(coefficients=coefficients), + stiffness=zeros, + damping=zeros, + ), + joint=JointParams( + stiffness=jnp.stack(joint_stiffness), + damping=jnp.stack(joint_damping), + ), + ) + young_operator, shear_operator, damping_operator = material_operators_from_params( + params, structure + ) + link_stiffness = [] + link_damping = [] + for index, segment in enumerate(input_segments): + link_dof = link_dofs[index] + if segment.link.stiffness is None: + stiffness = ( + segment.link.young_modulus * young_operator[index] + + segment.link.shear_modulus * shear_operator[index] + ) + else: + stiffness = _pad_matrix( + segment.link.stiffness, link_dof, layout_max_dof, "link stiffness" + ) + if segment.link.damping is None: + damping = ( + segment.link.material_damping_coefficient * damping_operator[index] + ) + else: + damping = _pad_matrix( + segment.link.damping, link_dof, layout_max_dof, "link damping" + ) + link_stiffness.append(stiffness) + link_damping.append(damping) + params = params.replace( + link=params.link.replace( + stiffness=jnp.stack(link_stiffness), damping=jnp.stack(link_damping) + ) + ) + params.validate_against_structure(structure) return params, structure diff --git a/src/soromox/systems/gvs/core.py b/src/soromox/systems/gvs/core.py index 60b69a0a..d09d4ee4 100644 --- a/src/soromox/systems/gvs/core.py +++ b/src/soromox/systems/gvs/core.py @@ -1,5 +1,5 @@ import math -from typing import Any +from typing import Any, Self import equinox as eqx import jax @@ -11,9 +11,17 @@ BaseThreadlikeRoutingParams, ThreadlikeRouting, ) +from soromox.systems.components import ( + ContinuumLinkParams, + CrossSectionGeometry, + IsotropicMaterialParams, +) from soromox.systems.gvs._assembly import assign_gvs_runtime_arrays from soromox.systems.gvs._runtime import SegmentRuntimeData -from soromox.systems.gvs.construction import params_and_structure_from_segments +from soromox.systems.gvs.construction import ( + material_operators_from_params, + params_and_structure_from_segments, +) from soromox.systems.gvs.joint_bases import ( B_Cylindrical, B_Fixed, @@ -25,7 +33,7 @@ B_Spherical, ) from soromox.systems.gvs.operands import GeometricOperand, JointOperand -from soromox.systems.gvs.params import GVSLinkParams, GVSParams +from soromox.systems.gvs.params import GVSParams from soromox.systems.gvs.primitives import Basis, Joint, Link from soromox.systems.gvs.specs import GVSSegment from soromox.systems.gvs.strain_bases import ( @@ -48,7 +56,7 @@ GVSStrainBasisStructure, GVSStructure, ) -from soromox.systems.soft_robot import CrossSectionGeometry, SoftRobot +from soromox.systems.soft_robot import SoftRobot from soromox.utils.geometry import poses from soromox.utils.integration import gauss_quadrature from soromox.utils.lie_algebra import constant_strain, se3, so3 @@ -79,16 +87,19 @@ class GVS(SoftRobot): num_integration_points: Number of integration/evaluation points for each link. dofs_per_segment: Number of DOFs for each link/joint pair (shape: num_segments x 2). integration_points, integration_weights: Gauss quadrature nodes and weights for each link. - mass_matrices, stiffness_matrices, damping_matrices: Mass, stiffness, and damping matrices at integration points. + mass_matrices: Mass matrices at integration points. Generalized link + stiffness and damping are stored canonically in ``params.link``. B_joint, B_Xs, B_Z1, B_Z2: Basis matrices for joints and links at quadrature points and intermediate points. xi_ref_joint, xi_ref_Xs, xi_ref_Z1, xi_ref_Z2: Reference strain vectors for joints and links. - joint_stiffness: Joint stiffness matrices. + joint_stiffness: Cached padded joint stiffness matrices. Canonical joint + stiffness and damping are stored in ``params.joint``. K_full: Precomputed full stiffness matrix for the robot. D_full: Precomputed full damping matrix for the robot. g0: Initial pose of the robot base as an SE(3) transformation matrix. g: Gravitational acceleration vector in 6D wrench form. basis_type_index: Index of the strain basis type used for each segment. - basis_active_params, basis_order_params: Parameters controlling the strain basis DOFs and orders. + basis_active_params, basis_order_params: Parameters controlling the strain + basis DOFs and orders. Notes ----- @@ -165,8 +176,6 @@ class GVS(SoftRobot): inner_integration_weights: Array mass_matrices: Array inner_mass_matrices: Array - stiffness_matrices: Array - damping_matrices: Array B_joint: Array B_Xs: Array @@ -184,6 +193,9 @@ class GVS(SoftRobot): K: Array # Active-coordinate stiffness matrix (num_dofs, num_dofs) D_full: Array # Full damping matrix (num_padded_dofs, num_padded_dofs) D_active: Array # Active-coordinate damping matrix (num_dofs, num_dofs) + young_stiffness_operator: Array + shear_stiffness_operator: Array + material_damping_operator: Array gather_indices: Array # Indices for active-to-padded coordinate gather gather_mask: Array # Valid-entry mask for active-to-padded coordinate gather @@ -213,7 +225,25 @@ def __init__( passive_elements: PassiveElement | tuple[PassiveElement, ...] | None = (), **kwargs: Any, ) -> None: - """Initialize a GVS robot from typed params and static segment structure.""" + """Initialize a GVS robot from typed parameters and static structure. + + Args: + params: Canonical dynamic link, joint, gravity, and base-pose + parameters. + structure: Static segment, strain-basis, quadrature, and padding + configuration corresponding to ``params``. + actuators: Optional actuator or tuple of actuators attached to the + model. + passive_elements: Optional passive element or tuple of passive + elements. Pass ``None`` or an empty tuple to disable them. + **kwargs: Additional keyword arguments forwarded to + :class:`BaseContinuumSoftRobot`. + + Raises: + TypeError: If ``params`` or ``structure`` has the wrong type. + ValueError: If the parameter arrays are inconsistent with the + static structure or an actuator configuration is invalid. + """ if not isinstance(params, GVSParams): raise TypeError("params must be a GVSParams instance.") if not isinstance(structure, GVSStructure): @@ -257,6 +287,29 @@ def params_from_segments( values are copied into ``GVSParams``. Omitted ``base_pose`` and ``gravity`` use the standard upright spatial mounting and negative-z Earth gravity. + + Args: + segments: Non-empty sequence of joint-link-basis segment + specifications. + gravity: Optional world-frame gravity vector with shape ``(3,)``. + base_pose: Optional base translation and unit quaternion with shape + ``(7,)``. + max_dof: Optional common padded dimension for every link and joint + generalized matrix. The required maximum is inferred when + omitted. + max_num_gauss_points: Optional common quadrature padding size. The + largest segment rule is used when omitted. + scale_rotational_basis_by_length: Whether to divide rotational + strain-basis rows by their link length. + + Returns: + A tuple containing canonical :class:`GVSParams` and its matching + static :class:`GVSStructure`. + + Raises: + KeyError: If a segment uses an unsupported joint or basis type. + ValueError: If no segments are supplied, a specification is + invalid, or a requested padding size is too small. """ return params_and_structure_from_segments( segments, @@ -283,6 +336,29 @@ def from_segments( Omitted ``base_pose`` and ``gravity`` use the standard upright spatial mounting and negative-z Earth gravity. + + Args: + segments: Non-empty sequence of joint-link-basis segment + specifications. + gravity: Optional world-frame gravity vector with shape ``(3,)``. + base_pose: Optional base translation and unit quaternion with shape + ``(7,)``. + max_dof: Optional common padded dimension for link and joint + generalized matrices. + max_num_gauss_points: Optional common quadrature padding size. + scale_rotational_basis_by_length: Whether to normalize rotational + strain-basis rows by link length. + **kwargs: Additional keyword arguments forwarded to the GVS + constructor, such as actuators or passive elements. + + Returns: + A fully initialized GVS model. + + Raises: + TypeError: If a constructor argument has an invalid type. + KeyError: If a segment uses an unsupported joint or basis type. + ValueError: If a segment specification, padding size, parameter + array, or attached component is invalid. """ params, structure = cls.params_from_segments( segments, @@ -296,16 +372,40 @@ def from_segments( @property def is_planar(self) -> bool: - """GVS is a spatial (3D) model.""" + """Return whether the system is planar. + + Returns: + Always ``False`` because GVS uses spatial rigid transformations. + """ return False @property def segment_length(self) -> Array: - """Per-segment backbone lengths.""" + """Return the per-segment backbone lengths. + + Returns: + Array with shape ``(num_segments,)``. + """ return jnp.asarray(self.segment_lengths) def cross_section_geometry(self, q: Array, s: Array) -> tuple[Array, Array]: - """Evaluate the configured circular, rectangular, or elliptical section.""" + """Evaluate cross-section geometry at a backbone coordinate. + + Args: + q: Generalized coordinates. Cross-section geometry is currently + configuration-independent, but ``q`` is accepted to implement + the common continuum-robot interface. + s: Scalar global backbone coordinate. + + Returns: + A tuple ``(geometry, dimensions)``. ``geometry`` is the integer + :class:`CrossSectionGeometry` tag. ``dimensions`` contains radius + for a circular section, width and height for a rectangle, or + semi-major and semi-minor radii for an ellipse. + + Raises: + ValueError: If ``s`` cannot be assigned to a segment. + """ segment_idx, s_local = self.classify_segment(s) length_i = self.segment_lengths[segment_idx] x = jnp.where(length_i > self.global_eps, s_local / length_i, 0.0) @@ -340,10 +440,7 @@ def _build_segment_i( max_num_integration_points: int, *, length: Array, - young_modulus: Array, - poisson_ratio: Array, density: Array, - damping_coefficient: Array, radius_initial: Array, radius_final: Array, height_initial: Array, @@ -442,10 +539,7 @@ def _build_segment_i( cross_section_geometry = link_structure.cross_section_geometry cross_section_geometry_idx = int(cross_section_geometry) - E = jnp.asarray(young_modulus) - nu = jnp.asarray(poisson_ratio) rho = jnp.asarray(density) - eta = jnp.asarray(damping_coefficient) L = jnp.asarray(length) r_i = jnp.asarray(radius_initial) @@ -459,8 +553,6 @@ def _build_segment_i( b_i = jnp.asarray(semi_minor_initial) b_f = jnp.asarray(semi_minor_final) - G = E / (2 * (1 + nu)) # Shear modulus - r_params = (r_i, r_f) h_params = (h_i, h_f) w_params = (w_i, w_f) @@ -470,8 +562,8 @@ def _build_segment_i( # === Basis attributes basetype = basis_structure.type basistype_idx = Basis.BASISTYPE_MAP[basetype] - Bdof = jnp.asarray(basis_structure.active).flatten() - Bodr = jnp.asarray(basis_structure.orders).flatten() + Bdof = jnp.asarray(basis_structure.strain_selector).flatten() + Bodr = jnp.asarray(basis_structure.basis_order).flatten() xi_ref = jnp.asarray(reference_strain).reshape(6, 1) dof_link = lax.switch( @@ -555,14 +647,7 @@ def xi_reffn(x: Array) -> Array: # Prepare the component vectors Ms_diag = jnp.stack([Ix_p, Iy_p, Iz_p, A_p, A_p, A_p], axis=1) # Shape: (np, 6) - Es_diag = jnp.stack( - [G * Ix_p, E * Iy_p, E * Iz_p, E * A_p, G * A_p, G * A_p], axis=1 - ) - Gs_diag = jnp.stack([Ix_p, 3 * Iy_p, 3 * Iz_p, 3 * A_p, A_p, A_p], axis=1) - Ms = rho * vmap(jnp.diag)(Ms_diag) # Shape: (np, 6, 6) - Es = vmap(jnp.diag)(Es_diag) - Gs = eta * vmap(jnp.diag)(Gs_diag) # Pad the arrays to the maximum number of integration points and DOFs integration_points_full = jnp.pad( @@ -585,25 +670,6 @@ def xi_reffn(x: Array) -> Array: ), mode="constant", ) - Es_full = jnp.pad( - Es, - ( - (0, max_num_integration_points - num_integration_points_i), - (0, 0), - (0, 0), - ), - mode="constant", - ) - Gs_full = jnp.pad( - Gs, - ( - (0, max_num_integration_points - num_integration_points_i), - (0, 0), - (0, 0), - ), - mode="constant", - ) - B_Xs_full = jnp.pad( B_Xs, ( @@ -662,8 +728,6 @@ def xi_reffn(x: Array) -> Array: integration_points=integration_points_full, integration_weights=integration_weights_full, mass_matrices=Ms_full, - stiffness_matrices=Es_full, - damping_matrices=Gs_full, B_joint=B_joint_full, B_Xs=B_Xs_full, B_Z1=B_Z1_full, @@ -676,10 +740,15 @@ def xi_reffn(x: Array) -> Array: ) def precompute(self) -> None: - """ - Precompute any necessary matrices or values for the simulation. + """Precompute padded gathers and canonical system matrices. - This method can be expanded to include additional precomputations as needed. + This refreshes material unit-response operators and projects the stored + link and joint stiffness and damping matrices into active generalized + coordinates. + + Returns: + ``None``. Cached arrays on the model are replaced in place during + initial construction. """ dofs_flat = self.dofs_per_segment.reshape(-1) start_indices_flat = jnp.cumsum(jnp.pad(dofs_flat, (1, 0)))[:-1] @@ -687,6 +756,9 @@ def precompute(self) -> None: gather_indices = start_indices[..., None] + jnp.arange(self.max_dof) gather_mask = jnp.arange(self.max_dof) < self.dofs_per_segment[..., None] + young_operator, shear_operator, damping_operator = ( + material_operators_from_params(self.params, self.structure) + ) K_full = self._stiffness_full_matrix() D_full = self._damping_full_matrix() object.__setattr__( @@ -709,6 +781,9 @@ def precompute(self) -> None: ) object.__setattr__(self, "gather_indices", gather_indices) object.__setattr__(self, "gather_mask", gather_mask) + object.__setattr__(self, "young_stiffness_operator", young_operator) + object.__setattr__(self, "shear_stiffness_operator", shear_operator) + object.__setattr__(self, "material_damping_operator", damping_operator) object.__setattr__(self, "K_full", K_full) object.__setattr__( self, "K", self.active_dof_map.T @ K_full @ self.active_dof_map @@ -719,7 +794,7 @@ def precompute(self) -> None: ) def _link_parameter_arrays( - self, link: GVSLinkParams + self, link: ContinuumLinkParams ) -> tuple[ Array, Array, @@ -730,18 +805,16 @@ def _link_parameter_arrays( Array, Array, Array, - Array, - Array, ]: """ Build link-derived GVS arrays from dynamic link params. This helper recomputes only quantities that depend on the physical link - parameters: length, cross-section geometry, local mass matrices, local - stiffness matrices, and local damping matrices. It deliberately leaves the - joint basis, link strain basis, reference strain, and active coordinate - layout unchanged. That makes it suitable for parameter updates that keep - the GVS discretization and generalized coordinates fixed. + parameters: length, cross-section geometry, and local mass matrices. It + deliberately leaves the joint basis, link strain basis, reference strain, + canonical generalized matrices, and active-coordinate layout unchanged. + That makes it suitable for geometry updates that keep the GVS + discretization and generalized coordinates fixed. Returns: A tuple containing: @@ -751,10 +824,6 @@ def _link_parameter_arrays( ``(self.num_segments + 1,)``. - ``mass_matrices``: local mass matrices at stored quadrature nodes, shape ``(self.num_segments, self.max_num_integration_points, 6, 6)``. - - ``stiffness_matrices``: local stiffness matrices at stored quadrature nodes, - shape ``(self.num_segments, self.max_num_integration_points, 6, 6)``. - - ``damping_matrices``: local damping matrices at stored quadrature nodes, shape - ``(self.num_segments, self.max_num_integration_points, 6, 6)``. - ``cross_section_geometry_index``: cross-section enum indices, shape ``(self.num_segments,)``. - ``radius_params``: circular radius interpolation parameters, shape @@ -780,8 +849,6 @@ def _link_parameter_arrays( segment_length_items = [] mass_matrix_items = [] - stiffness_matrix_items = [] - damping_matrix_items = [] cross_section_geometry_items = [] radius_param_items = [] height_param_items = [] @@ -794,31 +861,30 @@ def _link_parameter_arrays( segment_structure.link.cross_section_geometry ) - E = jnp.asarray(link.young_modulus[i_segment], dtype=jnp.float64) - nu = jnp.asarray(link.poisson_ratio[i_segment], dtype=jnp.float64) rho = jnp.asarray(link.density[i_segment], dtype=jnp.float64) - eta = jnp.asarray(link.damping_coefficient[i_segment], dtype=jnp.float64) L = jnp.asarray(link.length[i_segment], dtype=jnp.float64) - - r_params = ( - jnp.asarray(link.radius_initial[i_segment], dtype=jnp.float64), - jnp.asarray(link.radius_final[i_segment], dtype=jnp.float64), - ) - h_params = ( - jnp.asarray(link.height_initial[i_segment], dtype=jnp.float64), - jnp.asarray(link.height_final[i_segment], dtype=jnp.float64), - ) - w_params = ( - jnp.asarray(link.width_initial[i_segment], dtype=jnp.float64), - jnp.asarray(link.width_final[i_segment], dtype=jnp.float64), - ) - a_params = ( - jnp.asarray(link.semi_major_initial[i_segment], dtype=jnp.float64), - jnp.asarray(link.semi_major_final[i_segment], dtype=jnp.float64), + coefficients = link.cross_section.coefficients[i_segment] + pairs = [] + start = 0 + for profile_type, count in zip( + segment_structure.link.cross_section_profile_types, + segment_structure.link.cross_section_profile_parameter_counts, + strict=True, + ): + values = coefficients[start : start + count] + pairs.append( + (values[0], values[0]) + if profile_type == "constant" + else (values[0], values[1]) + ) + start += count + zero_pair = (jnp.array(0.0), jnp.array(0.0)) + r_params = pairs[0] if cross_section_geometry_idx == 0 else zero_pair + h_params, w_params = ( + pairs if cross_section_geometry_idx == 1 else (zero_pair, zero_pair) ) - b_params = ( - jnp.asarray(link.semi_minor_initial[i_segment], dtype=jnp.float64), - jnp.asarray(link.semi_minor_final[i_segment], dtype=jnp.float64), + a_params, b_params = ( + pairs if cross_section_geometry_idx == 2 else (zero_pair, zero_pair) ) geometric_operand = GeometricOperand( @@ -835,18 +901,10 @@ def _link_parameter_arrays( operand=geometric_operand, ) - G = E / (2 * (1 + nu)) Ms_diag = jnp.stack([Ix_p, Iy_p, Iz_p, A_p, A_p, A_p], axis=1) - Es_diag = jnp.stack( - [G * Ix_p, E * Iy_p, E * Iz_p, E * A_p, G * A_p, G * A_p], - axis=1, - ) - Gs_diag = jnp.stack([Ix_p, 3 * Iy_p, 3 * Iz_p, 3 * A_p, A_p, A_p], axis=1) segment_length_items.append(L) mass_matrix_items.append(rho * vmap(jnp.diag)(Ms_diag)) - stiffness_matrix_items.append(vmap(jnp.diag)(Es_diag)) - damping_matrix_items.append(eta * vmap(jnp.diag)(Gs_diag)) cross_section_geometry_items.append(cross_section_geometry_idx) radius_param_items.append(jnp.stack(r_params)) height_param_items.append(jnp.stack(h_params)) @@ -865,8 +923,6 @@ def _link_parameter_arrays( segment_lengths, segment_end_positions, jnp.stack(mass_matrix_items), - jnp.stack(stiffness_matrix_items), - jnp.stack(damping_matrix_items), jnp.asarray(cross_section_geometry_items, dtype=jnp.int32), jnp.stack(radius_param_items), jnp.stack(height_param_items), @@ -899,7 +955,25 @@ def _reference_strain_runtime_arrays( return xi_ref_joint, xi_ref_Xs, xi_ref_Z1, xi_ref_Z2 def with_params(self, params: GVSParams) -> "GVS": - """Return an updated copy with a full typed parameter object.""" + """Return a model copy using a complete parameter PyTree. + + Geometry-dependent mass data and material unit-response operators are + rebuilt. Canonical stiffness and damping are taken directly from + ``params`` and are not implicitly regenerated from material variables. + + Args: + params: Complete replacement parameter object compatible with this + model's static :class:`GVSStructure`. + + Returns: + A new GVS model with refreshed runtime arrays and projected system + matrices. The original model is unchanged. + + Raises: + TypeError: If ``params`` is not a :class:`GVSParams`. + ValueError: If a parameter shape or value is invalid or changes the + model's static padded layout. + """ if not isinstance(params, GVSParams): raise TypeError("params must be a GVSParams instance.") params.validate_against_structure(self.structure) @@ -909,9 +983,9 @@ def with_params(self, params: GVSParams) -> "GVS": xi_ref_Xs, xi_ref_Z1, xi_ref_Z2, - ) = self._reference_strain_runtime_arrays(params.reference_strain) + ) = self._reference_strain_runtime_arrays(params.link.reference_strain) - joint_stiffness = jnp.asarray(params.joint_stiffness, dtype=jnp.float64) + joint_stiffness = jnp.asarray(params.joint.stiffness, dtype=jnp.float64) if joint_stiffness.shape != self.joint_stiffness.shape: raise ValueError( "joint_stiffness shape changes the padded GVS layout; " @@ -931,8 +1005,6 @@ def with_params(self, params: GVSParams) -> "GVS": model.segment_lengths, model.segment_end_positions, model.mass_matrices, - model.stiffness_matrices, - model.damping_matrices, model.cross_section_geometry_index, model.radius_params, model.height_params, @@ -964,10 +1036,16 @@ def with_params(self, params: GVSParams) -> "GVS": ) K_full = updated_self._stiffness_full_matrix() D_full = updated_self._damping_full_matrix() + young_operator, shear_operator, damping_operator = ( + material_operators_from_params(params, self.structure) + ) return eqx.tree_at( lambda model: ( model.inner_integration_weights, model.inner_mass_matrices, + model.young_stiffness_operator, + model.shear_stiffness_operator, + model.material_damping_operator, model.K_full, model.K, model.D_full, @@ -982,6 +1060,9 @@ def with_params(self, params: GVSParams) -> "GVS": updated_self.mass_matrices[ :, 1 : updated_self.max_num_integration_points - 1 ], + young_operator, + shear_operator, + damping_operator, K_full, updated_self.active_dof_map.T @ K_full @ updated_self.active_dof_map, D_full, @@ -990,9 +1071,114 @@ def with_params(self, params: GVSParams) -> "GVS": ) def update_params(self, **updates: Array) -> "GVS": - """Return an updated copy with selected typed parameter fields replaced.""" + """Return a copy with selected top-level parameter fields replaced. + + Args: + **updates: Fields of :class:`GVSParams` to replace, typically + ``link``, ``joint``, ``gravity``, or ``base_pose``. + + Returns: + A new validated GVS model containing the replacements. + + Raises: + TypeError: If an unknown field is supplied or a replacement has an + invalid type. + ValueError: If the resulting parameters are incompatible with the + model's static structure. + """ return self.with_params(self.params.replace(**updates)) + def update_link_params(self, **updates: Any) -> "GVS": + """Return a copy with selected continuum-link fields replaced. + + Args: + **updates: Fields of :class:`ContinuumLinkParams` to replace, such + as ``length``, ``density``, ``reference_strain``, + ``cross_section``, ``stiffness``, or ``damping``. + + Returns: + A new validated GVS model with refreshed geometry-dependent caches. + + Raises: + TypeError: If an unknown link field is supplied. + ValueError: If the replacement is invalid or changes the padded + model layout. + """ + return self.with_params( + self.params.replace(link=self.params.link.replace(**updates)) + ) + + def update_joint_params(self, **updates: Any) -> "GVS": + """Return a copy with selected joint fields replaced. + + Args: + **updates: Fields of :class:`JointParams` to replace. Supported + public fields are ``stiffness`` and ``damping``. + + Returns: + A new validated GVS model whose assembled matrices include the + replacement joint contributions. + + Raises: + TypeError: If an unknown joint field is supplied. + ValueError: If a replacement matrix has an invalid value or padded + shape. + """ + return self.with_params( + self.params.replace(joint=self.params.joint.replace(**updates)) + ) + + def link_matrices_from_material( + self, material: IsotropicMaterialParams + ) -> tuple[Array, Array]: + """Map isotropic material variables to canonical GVS link matrices. + + Args: + material: Scalar or per-segment Young's modulus, shear modulus, and + material damping coefficient. Scalars are broadcast over all + segments. + + Returns: + A tuple ``(stiffness, damping)`` whose arrays both have shape + ``(num_segments, max_dof, max_dof)``. Inactive padded coordinates + remain zero. + + Raises: + ValueError: If a material field is not scalar or does not have + shape ``(num_segments,)``. + """ + material = material.broadcast(self.num_segments) + stiffness = ( + material.young_modulus[:, None, None] * self.young_stiffness_operator + + material.shear_modulus[:, None, None] * self.shear_stiffness_operator + ) + damping = ( + material.material_damping_coefficient[:, None, None] + * self.material_damping_operator + ) + return stiffness, damping + + def with_isotropic_material(self, material: IsotropicMaterialParams) -> Self: + """Return a copy whose link matrices are built from isotropic material. + + The supplied material PyTree remains caller-owned and is not stored on + the model. Reapplying this method is therefore explicit after geometry + or material updates. + + Args: + material: Scalar or per-segment isotropic material variables. + + Returns: + A new GVS model containing the generated canonical link stiffness + and damping matrices. + + Raises: + ValueError: If a material field cannot be broadcast to one value + per segment. + """ + stiffness, damping = self.link_matrices_from_material(material) + return self.update_link_params(stiffness=stiffness, damping=damping) + # Gathering functions ========================================================= @eqx.filter_jit def _min_size_gathered(self, vec_min_size_flat: Array) -> Array: @@ -2700,8 +2886,15 @@ def step(carry: Array, i: Array) -> tuple[tuple[Array, Array], None]: @eqx.filter_jit def jacobian_bodyframe(self, q: Array, s: Array) -> Array: - """ - Compute the Jacobian of the forward kinematics at a point s along the robot in the body frame. + """Compute the body-frame geometric Jacobian at one point. + + Args: + q: Active generalized coordinates with shape ``(num_dofs,)``. + s: Scalar global backbone coordinate in the interval from zero to + the total robot length. + + Returns: + Body-frame Jacobian with shape ``(6, num_dofs)``. """ _, J_local = self._jacobian_bodyframe_terms(q, s) return J_local @@ -2885,8 +3078,16 @@ def step( def jacobian_and_arc_length_derivative_bodyframe( self, q: Array, s: Array ) -> tuple[Array, Array]: - """ - Compute the body-frame Jacobian and its arc-length derivative at ``s``. + """Compute a body-frame Jacobian and its arc-length derivative. + + Args: + q: Active generalized coordinates with shape ``(num_dofs,)``. + s: Scalar global backbone coordinate. + + Returns: + A tuple ``(J, J_s)`` containing the body-frame Jacobian and its + derivative with respect to global arc length. Both arrays have + shape ``(6, num_dofs)``. """ _, J_body, Js_body, _ = ( self._jacobian_and_arc_length_derivative_bodyframe_terms(q, s) @@ -2895,8 +3096,14 @@ def jacobian_and_arc_length_derivative_bodyframe( @eqx.filter_jit def jacobian_arc_length_derivative_bodyframe(self, q: Array, s: Array) -> Array: - """ - Compute the arc-length derivative of the body-frame Jacobian at ``s``. + """Compute the body-frame Jacobian derivative with respect to arc length. + + Args: + q: Active generalized coordinates with shape ``(num_dofs,)``. + s: Scalar global backbone coordinate. + + Returns: + Arc-length derivative with shape ``(6, num_dofs)``. """ _, Js = self.jacobian_and_arc_length_derivative_bodyframe(q, s) return Js @@ -3127,8 +3334,16 @@ def jacobian_inertialframe(self, q: Array, s: Array) -> Array: def jacobian_and_arc_length_derivative_inertialframe( self, q: Array, s: Array ) -> tuple[Array, Array]: - """ - Compute the inertial-frame Jacobian and its arc-length derivative at ``s``. + """Compute an inertial-frame Jacobian and its arc-length derivative. + + Args: + q: Active generalized coordinates with shape ``(num_dofs,)``. + s: Scalar global backbone coordinate. + + Returns: + A tuple ``(J, J_s)`` containing the inertial-frame Jacobian and its + derivative with respect to global arc length. Both arrays have + shape ``(6, num_dofs)``. """ g_s, J_body, Js_body, gs = ( self._jacobian_and_arc_length_derivative_bodyframe_terms(q, s) @@ -3139,8 +3354,14 @@ def jacobian_and_arc_length_derivative_inertialframe( @eqx.filter_jit def jacobian_arc_length_derivative_inertialframe(self, q: Array, s: Array) -> Array: - """ - Compute the arc-length derivative of the inertial-frame Jacobian at ``s``. + """Compute the inertial Jacobian derivative with respect to arc length. + + Args: + q: Active generalized coordinates with shape ``(num_dofs,)``. + s: Scalar global backbone coordinate. + + Returns: + Arc-length derivative with shape ``(6, num_dofs)``. """ g_s, J_body, Js_body, gs = ( self._jacobian_and_arc_length_derivative_bodyframe_terms(q, s) @@ -3686,8 +3907,16 @@ def step( def jacobian_and_time_derivative_bodyframe( self, q: Array, qd: Array, s: Array ) -> tuple[Array, Array]: - """ - Compute the Jacobian time derivative of the forward kinematics at a point s along the robot in the body frame. + """Compute a body-frame Jacobian and its time derivative. + + Args: + q: Active generalized coordinates with shape ``(num_dofs,)``. + qd: Active generalized velocities with shape ``(num_dofs,)``. + s: Scalar global backbone coordinate. + + Returns: + A tuple ``(J, J_dot)`` containing the body-frame Jacobian and its + time derivative. Both arrays have shape ``(6, num_dofs)``. """ _, J_local, Jd_local = self._jacobian_and_time_derivative_bodyframe_terms( q, qd, s @@ -4413,81 +4642,13 @@ def point_term(j: Array) -> Array: @eqx.filter_jit def _stiffness_full_matrix(self) -> Array: - """ - Compute the full stiffness matrix of the robot. - - Returns: - K_full (Array): Full stiffness matrix, shape (num_segments * 2 * max_dof, num_segments * 2 * max_dof) - """ - - def K_i(i: Array) -> Array: - """Assemble stiffness contributions for one segment by quadrature. - - Args: - i (Array): Segment index (int). - - Returns: - K_blocks_i (Array): Two blocks (joint/link) of shape (2, max_dof, max_dof). - """ - # Joint ============================== - K_joint_i = jnp.zeros( - (self.max_dof, self.max_dof) - ) # self.joint_stiffness[i_segment] # (max_dof, max_dof) TODO - - # Link =============================== - length_i = self.segment_lengths[i] - Ws_i = self.integration_weights[i] # (max_num_integration_points, 1, ) - Es_i = self.stiffness_matrices[i] # (max_num_integration_points, 6, 6) - B_Xs_i = self.B_Xs[i] # (max_num_integration_points, 6, max_dof) - - def K_ij(j: Array) -> Array: - """ - Stiffness block at a single quadrature point. - - Args: - j (Array): Quadrature point index (int). - - Returns: - K_ij (Array): Block of shape (max_dof, max_dof). - """ - Ws_ij = Ws_i[j] - Es_ij = Es_i[j] # (6, 6) - B_Xs_ij = B_Xs_i[j] # (6, max_dof) - - if self.scale_rotational_basis_by_length: - B_Xs_ij = B_Xs_ij.at[:3, :].divide(length_i) - - return Ws_ij * (B_Xs_ij.T @ Es_ij @ B_Xs_ij) - - # we can skip the first and last quadrature points since their weight is zero - K_link_i = ( - jnp.sum( - vmap(K_ij)(jnp.arange(1, self.max_num_integration_points - 1)), - axis=0, - ) - * length_i - ) # (max_num_integration_points - 2, max_dof, max_dof) - - # Create a (2, max_dof, max_dof) array with K_joint_i and K_segment_i - K_blocks_i = jnp.stack([K_joint_i, K_link_i], axis=0) - return K_blocks_i - - K_blocks_tot = vmap(K_i)( - jnp.arange(self.num_segments) - ) # (num_segments, 2, max_dof, max_dof) - - # Assume that K_blocks is of the form (num_segments, 2, max_dof, max_dof) - K_blocks_flat = K_blocks_tot.reshape( - -1, self.max_dof, self.max_dof - ) # (num_segments * 2, max_dof, max_dof) - - # Convert to list of matrices - K_blocks_list = [K_blocks_flat[i] for i in range(K_blocks_flat.shape[0])] - - # Building the diagonal matrix in blocks - K_full = jax.scipy.linalg.block_diag(*K_blocks_list) - - return K_full + """Assemble interleaved canonical joint/link stiffness blocks.""" + blocks = jnp.stack( + [self.params.joint.stiffness, self.params.link.stiffness], axis=1 + ).reshape(-1, self.max_dof, self.max_dof) + return jax.scipy.linalg.block_diag( + *[blocks[index] for index in range(blocks.shape[0])] + ) @eqx.filter_jit def stiffness_matrix(self) -> Array: @@ -4514,80 +4675,13 @@ def elastic_force(self, q: Array) -> Array: @eqx.filter_jit def _damping_full_matrix(self) -> Array: - """ - Compute the full damping matrix of the robot. - - Returns: - D_full (Array): Full damping matrix, shape (num_segments * 2 * max_dof, num_segments * 2 * max_dof) - """ - - def D_i(i: Array) -> Array: - """Assemble damping contributions for one segment by quadrature. - - Args: - i (Array): Segment index (int). - - Returns: - D_blocks_i (Array): Two blocks (joint/link) of shape (2, max_dof, max_dof). - """ - # Joint ============================== - D_joint_i = jnp.zeros( - (self.max_dof, self.max_dof) - ) # Initialize joint stiffness matrix - - # Link =============================== - length_i = self.segment_lengths[i] - Ws_i = self.integration_weights[i] # (max_num_integration_points, 1, ) - Gs_i = self.damping_matrices[i] # (max_num_integration_points, 6, 6) - B_Xs_i = self.B_Xs[i] # (max_num_integration_points, 6, max_dof) - - def D_ij(j: Array) -> Array: - """Damping block at a single quadrature point. - - Args: - j (Array): Quadrature point index (int). - - Returns: - D_ij (Array): Block of shape (max_dof, max_dof). - """ - Ws_j = Ws_i[j] - Gs_j = Gs_i[j] # (6, 6) - B_Xs_j = B_Xs_i[j] # (6, max_dof) - - if self.scale_rotational_basis_by_length: - B_Xs_j = B_Xs_j.at[:3, :].divide(length_i) - - return Ws_j * (B_Xs_j.T @ Gs_j @ B_Xs_j) - - # we can skip the first and last quadrature points since their weight is zero - D_link_i = ( - jnp.sum( - vmap(D_ij)(jnp.arange(1, self.max_num_integration_points - 1)), - axis=0, - ) - * length_i - ) # (max_num_integration_points - 2, max_dof, max_dof) - - # Create a (2, max_dof, max_dof) array with D_joint_i and D_segment_i - D_blocks_i = jnp.stack([D_joint_i, D_link_i], axis=0) - return D_blocks_i - - D_blocks_tot = vmap(D_i)( - jnp.arange(self.num_segments) - ) # (num_segments, 2, max_dof, max_dof) - - # Assume that D_blocks is of the form (num_segments, 2, max_dof, max_dof) - D_blocks_flat = D_blocks_tot.reshape( - -1, self.max_dof, self.max_dof - ) # (num_segments * 2, max_dof, max_dof) - - # Convert to list of matrices - D_blocks_list = [D_blocks_flat[i] for i in range(D_blocks_flat.shape[0])] - - # Building the diagonal matrix in blocks - D_full = jax.scipy.linalg.block_diag(*D_blocks_list) - - return D_full + """Assemble interleaved canonical joint/link damping blocks.""" + blocks = jnp.stack( + [self.params.joint.damping, self.params.link.damping], axis=1 + ).reshape(-1, self.max_dof, self.max_dof) + return jax.scipy.linalg.block_diag( + *[blocks[index] for index in range(blocks.shape[0])] + ) @eqx.filter_jit def damping_matrix(self, q: Array) -> Array: diff --git a/src/soromox/systems/gvs/params.py b/src/soromox/systems/gvs/params.py index 53dbca8e..c0348744 100644 --- a/src/soromox/systems/gvs/params.py +++ b/src/soromox/systems/gvs/params.py @@ -1,191 +1,152 @@ -__all__ = ["GVSLinkParams", "GVSParams"] +"""Typed dynamic parameters for the geometric variable-strain system.""" -from typing import ClassVar +from __future__ import annotations -from jax import Array -from jax import numpy as jnp +from typing import TYPE_CHECKING, ClassVar -from soromox.systems.params import ( - BaseSoftRobotParams, - BaseSystemParams, - validate_quaternion_base_pose, -) +import jax.numpy as jnp +from soromox.systems.components import ContinuumLinkParams, JointParams +from soromox.systems.params import BaseSoftRobotParams, validate_quaternion_base_pose -class GVSLinkParams(BaseSystemParams): - """Dynamic per-link arrays for all GVS segments. +if TYPE_CHECKING: + from soromox.systems.gvs.structures import GVSStructure - Every field has leading shape ``(num_segments,)``; this is not a single-link - object. It stores the numeric link values without duplicating the static - cross-section family stored in ``GVSStructure.segments``. ``length`` follows - the singular per-link naming convention used by the other fields. Reference - strain is intentionally stored on ``GVSParams`` because it belongs to the - strain basis state, not the link cross-section/material data. - ``damping_coefficient`` is a viscosity-like modulus in Pa*s (N*s/m^2). - """ - - length: Array - young_modulus: Array - poisson_ratio: Array - density: Array - damping_coefficient: Array - radius_initial: Array - radius_final: Array - height_initial: Array - height_final: Array - width_initial: Array - width_final: Array - semi_major_initial: Array - semi_major_final: Array - semi_minor_initial: Array - semi_minor_final: Array - - def validate(self) -> None: - length = jnp.asarray(self.length) - if len(length.shape) != 1: - raise ValueError( - "length must be one-dimensional with shape (num_segments,)." - ) - n_segments = length.shape[0] - expected_shape = (n_segments,) - for name in ( - "young_modulus", - "poisson_ratio", - "density", - "damping_coefficient", - "radius_initial", - "radius_final", - "height_initial", - "height_final", - "width_initial", - "width_final", - "semi_major_initial", - "semi_major_final", - "semi_minor_initial", - "semi_minor_final", - ): - value = jnp.asarray(getattr(self, name)) - if value.shape != expected_shape: - raise ValueError( - f"{name} must have shape {expected_shape}, got {value.shape}." - ) +__all__ = ["GVSParams"] class GVSParams(BaseSoftRobotParams): - """Dynamic parameters for a GVS model. - - ``link`` contains per-segment link arrays. ``reference_strain`` has shape - ``(num_segments, 6)`` and is kept here, rather than in ``GVSLinkParams``, - because it parameterizes the strain basis/reference configuration rather - than link geometry or material properties. ``joint_stiffness`` has shape - ``(num_segments, max_dof, max_dof)`` and is padded to the static GVS layout. - ``base_pose`` uses scalar-first quaternion SE(3) coordinates - ``[qw, qx, qy, qz, x, y, z]`` with nonzero finite quaternion norm. Omitting - ``base_pose`` and ``gravity`` selects upright spatial mounting and - negative-z Earth gravity. + """Canonical dynamic parameters for a GVS model. + + Attributes: + link: Batched continuum-link parameters. Link stiffness and damping use + the padded generalized-coordinate dimension configured by the GVS + structure. + joint: Batched joint stiffness and damping matrices using the same + padded generalized-coordinate dimension as ``link``. + gravity: Gravity vector with shape ``(3,)`` in the world frame. + base_pose: Base translation and unit quaternion with shape ``(7,)``. """ is_planar: ClassVar[bool] = False - link: GVSLinkParams - reference_strain: Array - joint_stiffness: Array + link: ContinuumLinkParams + joint: JointParams def validate(self) -> None: + """Validate intrinsic shapes and values of the GVS parameters. + + Returns: + ``None`` after successful validation. + + Raises: + ValueError: If a component is invalid, link and joint batch shapes + disagree, reference strains do not have six components, or the + gravity and base-pose arrays have invalid shapes or values. + """ self.link.validate() - n_segments = self.link.length.shape[0] - gravity = jnp.asarray(self.gravity) - if gravity.shape != (3,): - raise ValueError(f"gravity must have shape (3,), got {gravity.shape}.") - validate_quaternion_base_pose("base_pose", self.base_pose, (7,)) - reference_strain = jnp.asarray(self.reference_strain) - if reference_strain.shape != (n_segments, 6): + self.joint.validate() + num_segments = self.link.length.shape[0] + if self.link.reference_strain.shape != (num_segments, 6): raise ValueError( - f"reference_strain must have shape ({n_segments}, 6), " - f"got {reference_strain.shape}." + "link.reference_strain must have shape " + f"({num_segments}, 6), got {self.link.reference_strain.shape}." ) - joint_stiffness = jnp.asarray(self.joint_stiffness) - if ( - joint_stiffness.ndim != 3 - or joint_stiffness.shape[0] != n_segments - or joint_stiffness.shape[1] != joint_stiffness.shape[2] - ): + if self.joint.stiffness.shape[0] != num_segments: + raise ValueError("joint params must contain one matrix per segment.") + if self.joint.stiffness.shape != self.link.stiffness.shape: raise ValueError( - "joint_stiffness must have shape " - f"({n_segments}, max_dof, max_dof), got {joint_stiffness.shape}." + "GVS link and joint matrices must use the same padded shape." ) + gravity = jnp.asarray(self.gravity) + if gravity.shape != (3,): + raise ValueError(f"gravity must have shape (3,), got {gravity.shape}.") + validate_quaternion_base_pose("base_pose", self.base_pose, (7,)) - def validate_against_structure(self, structure) -> None: - """Validate dynamic GVS arrays against the static padded layout.""" - from soromox.systems.gvs.primitives import Basis, Joint + def validate_against_structure(self, structure: GVSStructure) -> None: + """Validate dynamic arrays against a static padded GVS layout. + + Args: + structure: Static :class:`GVSStructure` defining the number of + segments, basis DOFs, cross-section coefficient packing, and + quadrature padding. + + Returns: + ``None`` after successful validation. + + Raises: + TypeError: If ``structure`` is not a :class:`GVSStructure`. + ValueError: If the structure is empty or any parameter batch, + padded matrix, cross-section coefficient array, basis DOF, or + quadrature setting is inconsistent with ``structure``. + """ + from soromox.systems.gvs.primitives import Joint from soromox.systems.gvs.structures import GVSStructure if not isinstance(structure, GVSStructure): raise TypeError("structure must be a GVSStructure instance.") if not structure.segments: raise ValueError("GVS requires at least one segment.") - self.validate() - n_segments = len(structure.segments) - if self.link.length.shape != (n_segments,): - raise ValueError( - f"link.length must have shape ({n_segments},), " - f"got {self.link.length.shape}." - ) - if self.reference_strain.shape != (n_segments, 6): - raise ValueError( - "reference_strain must have shape " - f"({n_segments}, 6), got {self.reference_strain.shape}." - ) - - max_num_gauss_points = structure.max_num_gauss_points - if max_num_gauss_points is not None and max_num_gauss_points < 5: + num_segments = len(structure.segments) + if self.link.length.shape != (num_segments,): + raise ValueError(f"link params must contain {num_segments} links.") + expected_coefficients = max( + sum(segment.link.cross_section_profile_parameter_counts) + for segment in structure.segments + ) + expected_cross_section_shape = (num_segments, expected_coefficients) + if self.link.cross_section.coefficients.shape != expected_cross_section_shape: raise ValueError( - f"max_num_gauss_points must be at least 5, got {max_num_gauss_points}." + "link.cross_section.coefficients must have shape " + f"{expected_cross_section_shape}, got " + f"{self.link.cross_section.coefficients.shape}." ) - dofs_joint = [ + joint_dofs = [ Joint.DICT_JOINT_TYPE_DOF[segment.joint.type] for segment in structure.segments ] - dofs_link = [ - int( - Basis.DOF_BRANCHES[Basis.BASISTYPE_MAP[segment.basis.type]]( - ( - jnp.asarray(segment.basis.active), - jnp.asarray(segment.basis.orders), + link_dofs = [] + for segment in structure.segments: + multiplier = 2 if segment.basis.type == "fourier" else 1 + offset = 1 + link_dofs.append( + sum( + int(active) * (multiplier * int(order) + offset) + for active, order in zip( + segment.basis.strain_selector, + segment.basis.basis_order, + strict=True, ) ) ) - for segment in structure.segments - ] - real_max_dof = max(dofs_joint + dofs_link) - max_dof = real_max_dof if structure.max_dof is None else structure.max_dof - if max_dof < real_max_dof: - raise ValueError( - "max_dof must be greater than or equal to the maximum DOF in links " - f"and joints ({real_max_dof}), but got {max_dof}." - ) - - expected_joint_stiffness_shape = (n_segments, max_dof, max_dof) - if self.joint_stiffness.shape != expected_joint_stiffness_shape: + required_max_dof = max(joint_dofs + link_dofs) + max_dof = required_max_dof if structure.max_dof is None else structure.max_dof + if max_dof < required_max_dof: raise ValueError( - "joint_stiffness must have shape " - f"{expected_joint_stiffness_shape}, got {self.joint_stiffness.shape}." + f"max_dof={max_dof} is smaller than required DOF {required_max_dof}." ) - - for idx, segment in enumerate(structure.segments): + expected = (num_segments, max_dof, max_dof) + for owner, value in ( + ("link.stiffness", self.link.stiffness), + ("link.damping", self.link.damping), + ("joint.stiffness", self.joint.stiffness), + ("joint.damping", self.joint.damping), + ): + if value.shape != expected: + raise ValueError( + f"{owner} must have shape {expected}, got {value.shape}." + ) + for index, segment in enumerate(structure.segments): if segment.num_gauss_points < 5: raise ValueError( - "Each GVS segment requires at least 5 Gauss points; " - f"segment {idx} has {segment.num_gauss_points}." + f"GVS segment {index} requires at least 5 Gauss points." ) if ( - max_num_gauss_points is not None - and max_num_gauss_points < segment.num_gauss_points + structure.max_num_gauss_points is not None + and structure.max_num_gauss_points < segment.num_gauss_points ): raise ValueError( - "max_num_gauss_points must be greater than or equal to the " - f"maximum segment num_gauss_points; segment {idx} has " - f"{segment.num_gauss_points}, got {max_num_gauss_points}." + "max_num_gauss_points must cover every segment quadrature rule." ) diff --git a/src/soromox/systems/gvs/specs.py b/src/soromox/systems/gvs/specs.py index 4d72f5fb..6d3ebde2 100644 --- a/src/soromox/systems/gvs/specs.py +++ b/src/soromox/systems/gvs/specs.py @@ -1,11 +1,15 @@ -from dataclasses import dataclass, field +"""GVS-specific segment and strain-basis construction specifications.""" + +from __future__ import annotations + +from dataclasses import dataclass from typing import Literal from jax import Array -from soromox.systems.soft_robot import CrossSectionGeometry +from soromox.systems.components import JointSpec, LinkSpec -__all__ = ["GVSSegment", "LinkSpec", "JointSpec", "StrainBasisSpec"] +__all__ = ["GVSSegment", "StrainBasisSpec"] StrainComponent = Literal[ "kappa_x", @@ -15,7 +19,6 @@ "sigma_y", "sigma_z", ] - BasisType = Literal[ "monomial", "legendre", @@ -25,18 +28,6 @@ "imq", ] -JointType = Literal[ - "revolute", - "prismatic", - "helical", - "cylindrical", - "planar", - "spherical", - "free", - "fixed", -] - - _COMPONENT_INDEX = { "kappa_x": 0, "kappa_y": 1, @@ -47,603 +38,72 @@ } -def _component_mask(active: Array | list | tuple[str, ...]) -> Array | list: - if isinstance(active, (list, tuple)) and all( - isinstance(item, str) for item in active +def _component_mask( + selector: Array | list | tuple[StrainComponent, ...], +) -> Array | list: + if isinstance(selector, (list, tuple)) and all( + isinstance(item, str) for item in selector ): mask = [0, 0, 0, 0, 0, 0] - for component in active: + for component in selector: try: mask[_COMPONENT_INDEX[component]] = 1 except KeyError as exc: - raise ValueError(f"Unknown strain component {component!r}") from exc + raise ValueError(f"Unknown strain component {component!r}.") from exc return mask - return active - - -def _orders_vector(orders: int | Array | list | tuple) -> Array | list: - if isinstance(orders, int): - return [orders] * 6 - return orders - - -@dataclass -class LinkSpec: - """Geometric and material specification of one GVS link. - - This is a single-link input object. When a GVS model is constructed, these - scalar values are packed into per-segment arrays in ``GVSParams.link``. - Initial and final cross-section dimensions are interpreted at normalized - arclength ``x = 0`` and ``x = 1`` and are linearly interpolated along the - link. - - Attributes: - cross_section_geometry: Cross-section family used to compute area and - second moments of area. Unit: dimensionless enum value. Shape: - scalar ``()``. - E: Young's modulus of the link material. Unit: pascal ``Pa``. Shape: - scalar ``()``. - nu: Poisson's ratio of the link material. Unit: dimensionless. Shape: - scalar ``()``. - rho: Mass density of the link material. Unit: ``kg/m^3``. Shape: - scalar ``()``. - eta: Material damping coefficient used to form the local damping - matrix. Unit: ``Pa*s`` (``N*s/m^2``). Shape: scalar ``()``. - L: Link arclength. Unit: meter ``m``. Shape: scalar ``()``. - r_i: Initial radius for circular cross-sections at ``x = 0``. Unit: - meter ``m``. Shape: scalar ``()``. Ignored for non-circular links. - r_f: Final radius for circular cross-sections at ``x = 1``. Unit: - meter ``m``. Shape: scalar ``()``. Ignored for non-circular links. - h_i: Initial height for rectangular cross-sections at ``x = 0``. Unit: - meter ``m``. Shape: scalar ``()``. Ignored for non-rectangular - links. - h_f: Final height for rectangular cross-sections at ``x = 1``. Unit: - meter ``m``. Shape: scalar ``()``. Ignored for non-rectangular - links. - w_i: Initial width for rectangular cross-sections at ``x = 0``. Unit: - meter ``m``. Shape: scalar ``()``. Ignored for non-rectangular - links. - w_f: Final width for rectangular cross-sections at ``x = 1``. Unit: - meter ``m``. Shape: scalar ``()``. Ignored for non-rectangular - links. - a_i: Initial semi-major axis for elliptical cross-sections at - ``x = 0``. Unit: meter ``m``. Shape: scalar ``()``. Ignored for - non-elliptical links. - a_f: Final semi-major axis for elliptical cross-sections at - ``x = 1``. Unit: meter ``m``. Shape: scalar ``()``. Ignored for - non-elliptical links. - b_i: Initial semi-minor axis for elliptical cross-sections at - ``x = 0``. Unit: meter ``m``. Shape: scalar ``()``. Ignored for - non-elliptical links. - b_f: Final semi-minor axis for elliptical cross-sections at - ``x = 1``. Unit: meter ``m``. Shape: scalar ``()``. Ignored for - non-elliptical links. - """ - - cross_section_geometry: CrossSectionGeometry - E: float - nu: float - rho: float - eta: float - L: float - r_i: float = 0.0 - r_f: float = 0.0 - h_i: float = 0.0 - h_f: float = 0.0 - w_i: float = 0.0 - w_f: float = 0.0 - a_i: float = 0.0 - a_f: float = 0.0 - b_i: float = 0.0 - b_f: float = 0.0 + return selector - @classmethod - def circular( - cls, - *, - E: float, - nu: float, - rho: float, - eta: float, - L: float, - r: float | None = None, - r_i: float | None = None, - r_f: float | None = None, - ) -> "LinkSpec": - """Create a circular-cross-section link specification. - Use ``r`` for a constant radius, or ``r_i`` and ``r_f`` for a radius - that varies linearly from the base to the tip. - - Args: - E: Young's modulus of the link material. Unit: pascal ``Pa``. - Shape: scalar ``()``. - nu: Poisson's ratio of the link material. Unit: dimensionless. - Shape: scalar ``()``. - rho: Mass density of the link material. Unit: ``kg/m^3``. Shape: - scalar ``()``. - eta: Material damping coefficient used to form the local damping - matrix. Unit: ``Pa*s`` (``N*s/m^2``). Shape: scalar ``()``. - L: Link arclength. Unit: meter ``m``. Shape: scalar ``()``. - r: Constant circular radius applied at both link ends. Unit: meter - ``m``. Shape: scalar ``()``. Optional if both ``r_i`` and - ``r_f`` are provided. - r_i: Initial circular radius at normalized arclength ``x = 0``. - Unit: meter ``m``. Shape: scalar ``()``. Optional if ``r`` is - provided. - r_f: Final circular radius at normalized arclength ``x = 1``. Unit: - meter ``m``. Shape: scalar ``()``. Optional if ``r`` is - provided. - - Returns: - LinkSpec: Single-link spec with ``cross_section_geometry`` set to - ``CrossSectionGeometry.CIRCULAR``. - - Raises: - ValueError: If neither ``r`` nor both endpoint radii are provided. - """ - if r is not None: - r_i = r if r_i is None else r_i - r_f = r if r_f is None else r_f - if r_i is None or r_f is None: - raise ValueError("Circular LinkSpec requires r or both r_i and r_f.") - return cls( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=E, - nu=nu, - rho=rho, - eta=eta, - L=L, - r_i=r_i, - r_f=r_f, - ) - - @classmethod - def rectangular( - cls, - *, - E: float, - nu: float, - rho: float, - eta: float, - L: float, - h: float | None = None, - w: float | None = None, - h_i: float | None = None, - h_f: float | None = None, - w_i: float | None = None, - w_f: float | None = None, - ) -> "LinkSpec": - """Create a rectangular-cross-section link specification. - - Use ``h`` and ``w`` for constant height and width, or use endpoint - values for dimensions that vary linearly from the base to the tip. - - Args: - E: Young's modulus of the link material. Unit: pascal ``Pa``. - Shape: scalar ``()``. - nu: Poisson's ratio of the link material. Unit: dimensionless. - Shape: scalar ``()``. - rho: Mass density of the link material. Unit: ``kg/m^3``. Shape: - scalar ``()``. - eta: Material damping coefficient used to form the local damping - matrix. Unit: ``Pa*s`` (``N*s/m^2``). Shape: scalar ``()``. - L: Link arclength. Unit: meter ``m``. Shape: scalar ``()``. - h: Constant rectangular height applied at both link ends. Unit: - meter ``m``. Shape: scalar ``()``. Optional if ``h_i`` and - ``h_f`` are provided. - w: Constant rectangular width applied at both link ends. Unit: - meter ``m``. Shape: scalar ``()``. Optional if ``w_i`` and - ``w_f`` are provided. - h_i: Initial rectangular height at normalized arclength ``x = 0``. - Unit: meter ``m``. Shape: scalar ``()``. Optional if ``h`` is - provided. - h_f: Final rectangular height at normalized arclength ``x = 1``. - Unit: meter ``m``. Shape: scalar ``()``. Optional if ``h`` is - provided. - w_i: Initial rectangular width at normalized arclength ``x = 0``. - Unit: meter ``m``. Shape: scalar ``()``. Optional if ``w`` is - provided. - w_f: Final rectangular width at normalized arclength ``x = 1``. - Unit: meter ``m``. Shape: scalar ``()``. Optional if ``w`` is - provided. - - Returns: - LinkSpec: Single-link spec with ``cross_section_geometry`` set to - ``CrossSectionGeometry.RECTANGULAR``. - - Raises: - ValueError: If neither constant dimensions nor all endpoint - dimensions are provided. - """ - if h is not None: - h_i = h if h_i is None else h_i - h_f = h if h_f is None else h_f - if w is not None: - w_i = w if w_i is None else w_i - w_f = w if w_f is None else w_f - if h_i is None or h_f is None or w_i is None or w_f is None: - raise ValueError( - "Rectangular LinkSpec requires h/w or h_i, h_f, w_i, and w_f." - ) - return cls( - cross_section_geometry=CrossSectionGeometry.RECTANGULAR, - E=E, - nu=nu, - rho=rho, - eta=eta, - L=L, - h_i=h_i, - h_f=h_f, - w_i=w_i, - w_f=w_f, - ) - - @classmethod - def elliptical( - cls, - *, - E: float, - nu: float, - rho: float, - eta: float, - L: float, - a: float | None = None, - b: float | None = None, - a_i: float | None = None, - a_f: float | None = None, - b_i: float | None = None, - b_f: float | None = None, - ) -> "LinkSpec": - """Create an elliptical-cross-section link specification. - - Use ``a`` and ``b`` for constant semi-major and semi-minor axes, or use - endpoint values for axes that vary linearly from the base to the tip. - - Args: - E: Young's modulus of the link material. Unit: pascal ``Pa``. - Shape: scalar ``()``. - nu: Poisson's ratio of the link material. Unit: dimensionless. - Shape: scalar ``()``. - rho: Mass density of the link material. Unit: ``kg/m^3``. Shape: - scalar ``()``. - eta: Material damping coefficient used to form the local damping - matrix. Unit: ``Pa*s`` (``N*s/m^2``). Shape: scalar ``()``. - L: Link arclength. Unit: meter ``m``. Shape: scalar ``()``. - a: Constant semi-major axis applied at both link ends. Unit: meter - ``m``. Shape: scalar ``()``. Optional if ``a_i`` and ``a_f`` - are provided. - b: Constant semi-minor axis applied at both link ends. Unit: meter - ``m``. Shape: scalar ``()``. Optional if ``b_i`` and ``b_f`` - are provided. - a_i: Initial semi-major axis at normalized arclength ``x = 0``. - Unit: meter ``m``. Shape: scalar ``()``. Optional if ``a`` is - provided. - a_f: Final semi-major axis at normalized arclength ``x = 1``. Unit: - meter ``m``. Shape: scalar ``()``. Optional if ``a`` is - provided. - b_i: Initial semi-minor axis at normalized arclength ``x = 0``. - Unit: meter ``m``. Shape: scalar ``()``. Optional if ``b`` is - provided. - b_f: Final semi-minor axis at normalized arclength ``x = 1``. Unit: - meter ``m``. Shape: scalar ``()``. Optional if ``b`` is - provided. - - Returns: - LinkSpec: Single-link spec with ``cross_section_geometry`` set to - ``CrossSectionGeometry.ELLIPTICAL``. - - Raises: - ValueError: If neither constant axes nor all endpoint axes are - provided. - """ - if a is not None: - a_i = a if a_i is None else a_i - a_f = a if a_f is None else a_f - if b is not None: - b_i = b if b_i is None else b_i - b_f = b if b_f is None else b_f - if a_i is None or a_f is None or b_i is None or b_f is None: - raise ValueError( - "Elliptical LinkSpec requires a/b or a_i, a_f, b_i, and b_f." - ) - return cls( - cross_section_geometry=CrossSectionGeometry.ELLIPTICAL, - E=E, - nu=nu, - rho=rho, - eta=eta, - L=L, - a_i=a_i, - a_f=a_f, - b_i=b_i, - b_f=b_f, - ) - - -@dataclass -class JointSpec: - """Kinematic and stiffness specification of the joint preceding a GVS link. - - The joint contributes the rigid joint coordinates before the soft link in a - segment. Joint strain bases use the spatial twist ordering - ``[omega_x, omega_y, omega_z, v_x, v_y, v_z]``. - - Attributes: - type: Joint family name. Unit: dimensionless categorical value. Shape: - scalar ``()``. Supported values are ``"revolute"``, - ``"prismatic"``, ``"helical"``, ``"cylindrical"``, ``"planar"``, - ``"spherical"``, ``"free"``, and ``"fixed"``. - axis: Axis used by revolute, prismatic, helical, and cylindrical - joints. Unit: dimensionless categorical value. Shape: scalar ``()``. - Valid values are ``"x"``, ``"y"``, and ``"z"``. - plane: Motion plane used by planar joints. Unit: dimensionless - categorical value. Shape: scalar ``()``. Valid values are ``"xy"``, - ``"yz"``, and ``"xz"``. - pitch: Translation per unit angular coordinate for helical joints. Unit: - meter per radian ``m/rad``. Shape: scalar ``()``. Ignored by - non-helical joints. - stiffness: Joint stiffness matrix for the active joint coordinates. - Unit: generalized joint effort per generalized joint coordinate - (for example ``N*m/rad`` for rotational coordinates and ``N/m`` for - translational coordinates). Shape: ``(joint_dof, joint_dof)``, - where ``joint_dof`` is 0 for fixed, 1 for revolute/prismatic/ - helical, 2 for cylindrical, 3 for planar/spherical, and 6 for free - joints. An empty list is accepted and interpreted as zero stiffness. - """ - - type: JointType - axis: Literal["x", "y", "z"] = "x" - plane: Literal["xy", "yz", "xz"] = "xy" - pitch: float = 0.0 - stiffness: Array | list = field(default_factory=list) - - @classmethod - def fixed(cls) -> "JointSpec": - """Create a fixed joint specification. - - The joint has no generalized coordinates and contributes no joint - motion before the link. - - Returns: - JointSpec: Fixed-joint spec with zero joint DOFs. The implied - stiffness has unit not applicable and shape ``(0, 0)``. - """ - return cls(type="fixed") - - @classmethod - def revolute( - cls, - axis: Literal["x", "y", "z"] = "x", - stiffness: Array | list | None = None, - ) -> "JointSpec": - """Create a revolute joint specification. - - Args: - axis: Axis of rotation. Unit: dimensionless categorical value. - Shape: scalar ``()``. Valid values are ``"x"``, ``"y"``, and - ``"z"``. - stiffness: Rotational stiffness matrix for the revolute coordinate. - Unit: ``N*m/rad``. Shape: ``(1, 1)``. If omitted or ``None``, - zero stiffness is used. - - Returns: - JointSpec: Revolute-joint spec with one angular coordinate in - radians about ``axis``. - """ - return cls( - type="revolute", axis=axis, stiffness=[] if stiffness is None else stiffness - ) - - @classmethod - def prismatic( - cls, - axis: Literal["x", "y", "z"] = "x", - stiffness: Array | list | None = None, - ) -> "JointSpec": - """Create a prismatic joint specification. - - Args: - axis: Axis of translation. Unit: dimensionless categorical value. - Shape: scalar ``()``. Valid values are ``"x"``, ``"y"``, and - ``"z"``. - stiffness: Translational stiffness matrix for the prismatic - coordinate. Unit: ``N/m``. Shape: ``(1, 1)``. If omitted or - ``None``, zero stiffness is used. - - Returns: - JointSpec: Prismatic-joint spec with one translational coordinate - in meters along ``axis``. - """ - return cls( - type="prismatic", - axis=axis, - stiffness=[] if stiffness is None else stiffness, - ) - - @classmethod - def helical( - cls, - axis: Literal["x", "y", "z"] = "x", - pitch: float = 0.0, - stiffness: Array | list | None = None, - ) -> "JointSpec": - """Create a helical joint specification. - - Args: - axis: Screw axis of the helical joint. Unit: dimensionless - categorical value. Shape: scalar ``()``. Valid values are - ``"x"``, ``"y"``, and ``"z"``. - pitch: Translation per unit angular coordinate along ``axis``. - Unit: meter per radian ``m/rad``. Shape: scalar ``()``. - stiffness: Generalized stiffness matrix for the helical coordinate. - Unit: generalized screw effort per radian. Shape: ``(1, 1)``. - If omitted or ``None``, zero stiffness is used. - - Returns: - JointSpec: Helical-joint spec with one screw coordinate in radians. - """ - return cls( - type="helical", - axis=axis, - pitch=pitch, - stiffness=[] if stiffness is None else stiffness, - ) - - @classmethod - def cylindrical( - cls, - axis: Literal["x", "y", "z"] = "x", - stiffness: Array | list | None = None, - ) -> "JointSpec": - """Create a cylindrical joint specification. - - Args: - axis: Shared rotation and translation axis. Unit: dimensionless - categorical value. Shape: scalar ``()``. Valid values are - ``"x"``, ``"y"``, and ``"z"``. - stiffness: Generalized stiffness matrix for coordinates - ``[rotation_about_axis, translation_along_axis]``. Units are - mixed by coordinate, with diagonal examples ``N*m/rad`` and - ``N/m``. Shape: ``(2, 2)``. If omitted or ``None``, zero - stiffness is used. - - Returns: - JointSpec: Cylindrical-joint spec with one angular coordinate in - radians and one translational coordinate in meters. - """ - return cls( - type="cylindrical", - axis=axis, - stiffness=[] if stiffness is None else stiffness, - ) - - @classmethod - def planar( - cls, - plane: Literal["xy", "yz", "xz"] = "xy", - stiffness: Array | list | None = None, - ) -> "JointSpec": - """Create a planar joint specification. - - Args: - plane: Plane in which the joint translates and about whose normal it - rotates. Unit: dimensionless categorical value. Shape: scalar - ``()``. Valid values are ``"xy"``, ``"yz"``, and ``"xz"``. - stiffness: Generalized stiffness matrix for coordinates - ``[rotation_normal_to_plane, translation_1, translation_2]``. - Units are mixed by coordinate, with diagonal examples - ``N*m/rad`` and ``N/m``. Shape: ``(3, 3)``. If omitted or - ``None``, zero stiffness is used. - - Returns: - JointSpec: Planar-joint spec with one angular coordinate in radians - and two translational coordinates in meters. - """ - return cls( - type="planar", - plane=plane, - stiffness=[] if stiffness is None else stiffness, - ) - - @classmethod - def spherical(cls, stiffness: Array | list | None = None) -> "JointSpec": - """Create a spherical joint specification. - - Args: - stiffness: Rotational stiffness matrix for the three angular - coordinates ``[rotation_x, rotation_y, rotation_z]``. Unit: - ``N*m/rad``. Shape: ``(3, 3)``. If omitted or ``None``, zero - stiffness is used. - - Returns: - JointSpec: Spherical-joint spec with three angular coordinates in - radians. - """ - return cls(type="spherical", stiffness=[] if stiffness is None else stiffness) - - @classmethod - def free(cls, stiffness: Array | list | None = None) -> "JointSpec": - """Create a free joint specification. - - Args: - stiffness: Generalized stiffness matrix for coordinates - ``[rotation_x, rotation_y, rotation_z, translation_x, - translation_y, translation_z]``. Units are mixed by coordinate, - with diagonal examples ``N*m/rad`` and ``N/m``. Shape: - ``(6, 6)``. If omitted or ``None``, zero stiffness is used. - - Returns: - JointSpec: Free-joint spec with three angular coordinates in - radians and three translational coordinates in meters. - """ - return cls(type="free", stiffness=[] if stiffness is None else stiffness) - - def __post_init__(self) -> None: - self.type = self.type.lower() +def _orders_vector(order: int | Array | list | tuple) -> Array | list: + return [order] * 6 if isinstance(order, int) else order @dataclass class StrainBasisSpec: - """Strain-basis parametrization for one GVS link. + """Describe the variable-strain basis used by one GVS link. - The basis maps link generalized coordinates to the six spatial strain - components ordered as ``[kappa_x, kappa_y, kappa_z, sigma_x, sigma_y, - sigma_z]``. + String-valued strain selectors are converted to the six-component mask used + internally by GVS. A scalar basis order is replicated for all six strain + components, including inactive components. Attributes: - type: Strain basis family name. Unit: dimensionless categorical value. - Shape: scalar ``()``. Supported values are ``"monomial"``, - ``"legendre"``, ``"chebyshev"``, ``"fourier"``, ``"gaussian"``, - and ``"imq"``. - active: Active strain-component selector. Unit: dimensionless mask. - Shape: ``(6,)`` after initialization. It may be passed either as a - six-entry mask in the strain order above or as a tuple/list of - component names such as ``("kappa_y", "sigma_x")``. - orders: Basis order for each strain component. Unit: dimensionless - integer order. Shape: scalar ``()`` when passed as an ``int`` (then - broadcast to all six components) or ``(6,)`` when passed as an - array/list/tuple. Stored shape after initialization is ``(6,)``. - xi_ref: Reference strain vector in the local link frame, ordered as - ``[kappa_x, kappa_y, kappa_z, sigma_x, sigma_y, sigma_z]``. Units: - angular strain entries ``kappa_*`` are ``rad/m``; translational - strain entries ``sigma_*`` are dimensionless stretch/shear values, - with ``sigma_x = 1`` representing a straight unstretched link along - the local x-axis. Shape: ``(6,)``. + type: Basis family. Supported values are ``"monomial"``, ``"legendre"``, + ``"chebyshev"``, ``"fourier"``, ``"gaussian"``, and ``"imq"``. + strain_selector: Active strain components, supplied either as a + six-element numeric mask or as a sequence of component names. The + supported names are ``kappa_x``, ``kappa_y``, ``kappa_z``, + ``sigma_x``, ``sigma_y``, and ``sigma_z``. + basis_order: Basis order supplied as one integer for every strain + component or as a six-element array-like value. """ type: BasisType - active: Array | list | tuple[StrainComponent, ...] - orders: int | Array | list | tuple - xi_ref: Array | list = field(default_factory=lambda: [0.0, 0.0, 0.0, 1.0, 0.0, 0.0]) + strain_selector: Array | list | tuple[StrainComponent, ...] + basis_order: int | Array | list | tuple def __post_init__(self) -> None: - self.type = self.type.lower() - self.active = _component_mask(self.active) - self.orders = _orders_vector(self.orders) + self.type = self.type.lower() # type: ignore[assignment] + self.strain_selector = _component_mask(self.strain_selector) + self.basis_order = _orders_vector(self.basis_order) @dataclass class GVSSegment: - """Complete user-facing specification of one GVS segment. + """Bundle all construction specifications for one GVS segment. - A segment consists of one joint followed by one soft GVS link. The segment - spec is the public construction unit consumed by ``GVS.from_segments`` and - ``GVS.params_from_segments``. + A segment consists of a discrete joint followed by a variable-strain + continuum link. The object is a construction specification only; calling + :meth:`GVS.from_segments ` converts + its values to canonical runtime parameters and static GVS structure. Attributes: - link: Link geometry, material, and length specification. Unit: not - applicable; nested ``LinkSpec`` object. Shape: scalar object ``()``. - joint: Specification of the joint preceding the link. Unit: not - applicable; nested ``JointSpec`` object. Shape: scalar object ``()``. - basis: Link strain-basis and reference-strain specification. Unit: not - applicable; nested ``StrainBasisSpec`` object. Shape: scalar object - ``()``. - num_gauss_points: Number of interior Gauss-Legendre quadrature points - requested for this link. Unit: points. Shape: scalar ``()``. The - runtime integration arrays include two additional zero-weight - boundary nodes, so their unpadded length is - ``num_gauss_points + 2``. Expected value is at least 5. + link: Continuum-link geometry, reference strain, and either isotropic + material properties or explicit generalized link matrices. + joint: Joint type, kinematics, and generalized joint stiffness and + damping in active joint coordinates. + basis: Variable-strain basis used to parameterize the link strain. + num_gauss_points: Number of Gauss points used for link quadrature. GVS + requires at least five points per segment. """ link: LinkSpec diff --git a/src/soromox/systems/gvs/structures.py b/src/soromox/systems/gvs/structures.py index 402d2540..a1b706c8 100644 --- a/src/soromox/systems/gvs/structures.py +++ b/src/soromox/systems/gvs/structures.py @@ -10,7 +10,7 @@ import equinox as eqx -from soromox.systems.soft_robot import CrossSectionGeometry +from soromox.systems.components import CrossSectionGeometry, ProfileType class GVSLinkStructure(eqx.Module): @@ -19,15 +19,29 @@ class GVSLinkStructure(eqx.Module): This stores only choices that affect compilation or dispatch. Numeric link values such as length, material constants, and cross-section dimensions live in ``GVSParams.link``. + + Attributes: + cross_section_geometry: Static solid cross-section family. + cross_section_profile_types: Profile type for each geometry dimension. + cross_section_profile_parameter_counts: Number of packed coefficients + consumed by each geometry dimension. """ cross_section_geometry: CrossSectionGeometry = eqx.field(static=True) + cross_section_profile_types: tuple[ProfileType, ...] = eqx.field(static=True) + cross_section_profile_parameter_counts: tuple[int, ...] = eqx.field(static=True) class GVSJointStructure(eqx.Module): """Static joint choices for one GVS segment. - Joint stiffness is dynamic and lives in ``GVSParams.joint_stiffness``. + Joint stiffness and damping are dynamic and live in ``GVSParams.joint``. + + Attributes: + type: Joint family name. + axis: Material-frame axis for axis-based joints. + plane: Motion plane for planar joints. + pitch: Translation per radian for a helical joint. """ type: str = eqx.field(static=True) @@ -39,16 +53,28 @@ class GVSJointStructure(eqx.Module): class GVSStrainBasisStructure(eqx.Module): """Static strain-basis choices for one GVS segment. - Reference strain is dynamic and lives in ``GVSParams.reference_strain``. + Reference strain is dynamic and lives in ``GVSParams.link.reference_strain``. + + Attributes: + type: Variable-strain basis family. + strain_selector: Six-element active-strain mask. + basis_order: Six-element basis-order tuple. """ type: str = eqx.field(static=True) - active: tuple[Any, ...] = eqx.field(static=True) - orders: tuple[Any, ...] = eqx.field(static=True) + strain_selector: tuple[Any, ...] = eqx.field(static=True) + basis_order: tuple[Any, ...] = eqx.field(static=True) class GVSSegmentStructure(eqx.Module): - """Static structure for one GVS segment.""" + """Static structure for one GVS segment. + + Attributes: + link: Static link geometry and profile choices. + joint: Static joint kinematic choices. + basis: Static link strain-basis choices. + num_gauss_points: Number of Gauss points used for segment quadrature. + """ link: GVSLinkStructure = eqx.field(static=True) joint: GVSJointStructure = eqx.field(static=True) @@ -62,6 +88,15 @@ class GVSStructure(eqx.Module): ``segments`` stores stripped static segment structures: joint families, basis families/orders/active masks, quadrature counts, and cross-section families. Dynamic numeric values live in ``GVSParams``. + + Attributes: + segments: Non-empty tuple of static segment structures. + max_dof: Common padded link and joint generalized-coordinate dimension, + or ``None`` before construction resolves it. + max_num_gauss_points: Common quadrature padding size, or ``None`` before + construction resolves it. + scale_rotational_basis_by_length: Whether rotational strain-basis rows + are divided by link length. """ segments: tuple[GVSSegmentStructure, ...] = eqx.field(static=True) diff --git a/src/soromox/systems/hsa/planar_hsa.py b/src/soromox/systems/hsa/planar_hsa.py index 8e3e420a..c1087c97 100644 --- a/src/soromox/systems/hsa/planar_hsa.py +++ b/src/soromox/systems/hsa/planar_hsa.py @@ -8,9 +8,10 @@ from jax import Array, jacfwd, lax from jax import numpy as jnp +from soromox.systems.components import CrossSectionGeometry from soromox.systems.hsa.params import PlanarHSAParams from soromox.systems.hsa.structures import PlanarHSAStructure -from soromox.systems.soft_robot import CrossSectionGeometry, SoftRobot +from soromox.systems.soft_robot import SoftRobot from soromox.utils.basic import ( compute_strain_basis, concatenate_params_syms, diff --git a/src/soromox/systems/params.py b/src/soromox/systems/params.py index 3f5444ca..9c1fe518 100644 --- a/src/soromox/systems/params.py +++ b/src/soromox/systems/params.py @@ -78,6 +78,9 @@ def validate_quaternion_base_pose( performed; transform helpers still avoid zero-norm division to keep traced code finite. + Returns: + ``None`` after successful validation. + Raises: ValueError: If the shape is wrong, any entry is non-finite, or the quaternion component is zero or numerically too small to normalize @@ -99,7 +102,21 @@ class BaseSystemParams(eqx.Module): """Base class for dynamic system parameters stored as JAX PyTrees.""" def replace(self, **updates: Any) -> "BaseSystemParams": - """Return a copy with selected fields replaced.""" + """Return an immutable copy with selected fields replaced. + + Args: + **updates: Parameter field names and their replacement values. + Nested parameter components should first be replaced with their + own ``replace`` method, then supplied as a top-level field. + + Returns: + A validated parameter PyTree of the same concrete class. + + Raises: + KeyError: If any update names an unknown parameter field. + TypeError: If a replacement cannot be inserted into the PyTree. + ValueError: If validation of the resulting parameters fails. + """ valid_names = {field.name for field in fields(self)} unknown_names = set(updates) - valid_names if unknown_names: @@ -123,11 +140,27 @@ def _normalize_replacement(self, name: str, value: Any) -> Any: return value def validate(self) -> None: - """Validate parameter consistency.""" + """Validate intrinsic parameter consistency. + + Returns: + ``None``. Subclasses override this method to raise on invalid + shapes, values, or component relationships. + """ return None def validate_against_structure(self, structure: Any) -> None: - """Validate params against static construction choices.""" + """Validate parameters against static construction choices. + + Args: + structure: System-specific static structure to validate against. + + Returns: + ``None`` after successful validation. + + Raises: + ValueError: If intrinsic or structure-dependent validation fails in + a subclass implementation. + """ self.validate() @@ -230,21 +263,70 @@ def _from_mounting( def horizontal( cls, *, base_position: Array | None = None, **kwargs: Any ) -> "BaseSoftRobotParams": - """Construct parameters with the backbone pointing along world +x.""" + """Construct parameters with the backbone pointing along world +x. + + Args: + base_position: Optional base translation with shape ``(2,)`` for a + planar model or ``(3,)`` for a spatial model. + **kwargs: Remaining arguments forwarded to the concrete parameter + class constructor. ``base_pose`` is not accepted. + + Returns: + A concrete parameter object in the horizontal mounting. + + Raises: + TypeError: If the concrete class has not declared its planarity or + ``base_pose`` is supplied. + ValueError: If ``base_position`` has an invalid shape or value. + """ return cls._from_mounting("horizontal", base_position=base_position, **kwargs) @classmethod def upright( cls, *, base_position: Array | None = None, **kwargs: Any ) -> "BaseSoftRobotParams": - """Construct parameters pointing along world +y (planar) or +z (spatial).""" + """Construct parameters pointing along world +y or +z. + + Planar models point along world +y; spatial models point along world +z. + + Args: + base_position: Optional base translation with shape ``(2,)`` for a + planar model or ``(3,)`` for a spatial model. + **kwargs: Remaining arguments forwarded to the concrete parameter + class constructor. ``base_pose`` is not accepted. + + Returns: + A concrete parameter object in the upright mounting. + + Raises: + TypeError: If the concrete class has not declared its planarity or + ``base_pose`` is supplied. + ValueError: If ``base_position`` has an invalid shape or value. + """ return cls._from_mounting("upright", base_position=base_position, **kwargs) @classmethod def hanging( cls, *, base_position: Array | None = None, **kwargs: Any ) -> "BaseSoftRobotParams": - """Construct parameters pointing along world -y (planar) or -z (spatial).""" + """Construct parameters pointing along world -y or -z. + + Planar models point along world -y; spatial models point along world -z. + + Args: + base_position: Optional base translation with shape ``(2,)`` for a + planar model or ``(3,)`` for a spatial model. + **kwargs: Remaining arguments forwarded to the concrete parameter + class constructor. ``base_pose`` is not accepted. + + Returns: + A concrete parameter object in the hanging mounting. + + Raises: + TypeError: If the concrete class has not declared its planarity or + ``base_pose`` is supplied. + ValueError: If ``base_position`` has an invalid shape or value. + """ return cls._from_mounting("hanging", base_position=base_position, **kwargs) def _normalize_replacement(self, name: str, value: Any) -> Any: @@ -258,15 +340,12 @@ def _normalize_replacement(self, name: str, value: Any) -> Any: class BaseContinuumSoftRobotParams(BaseSoftRobotParams): """Shared dynamic parameters for continuum soft robots. - Field names denote one segment's physical quantity; the leading axis stores - the segment batch. ``reference_strain`` contains the flattened per-segment - reference strain used by PCS-style continuum models. + Concrete continuum systems expose their batched physical values through a + shared ``ContinuumLinkParams`` field named ``link``. The field is declared + by concrete subclasses to keep this base module independent from the + reusable component package. """ - length: Array - density: Array - reference_strain: Array - class BaseArticulatedSoftRobotParams(BaseSoftRobotParams): """Shared dynamic parameters for articulated systems. diff --git a/src/soromox/systems/pcs/isupport.py b/src/soromox/systems/pcs/isupport.py index 11c9dd25..170f399f 100644 --- a/src/soromox/systems/pcs/isupport.py +++ b/src/soromox/systems/pcs/isupport.py @@ -6,9 +6,9 @@ from jax import Array, vmap from soromox.actuation import ThreadlikeActuator, ThreadlikeRouting +from soromox.systems.components import ContinuumLinkParams, CrossSectionParams from soromox.systems.pcs.params import ISupportParams from soromox.systems.pcs.structures import ISupportStructure, PCSStructure -from soromox.utils.array_math import blk_diag from .pcs import PCS @@ -91,7 +91,7 @@ def _normalize_structure_for_params( params: ISupportParams, structure: ISupportStructure, ) -> ISupportStructure: - num_physical_segments = int(params.length.shape[0]) + num_physical_segments = int(params.link.length.shape[0]) rigid_segment_selector = structure.rigid_segment_selector if rigid_segment_selector is None: @@ -140,7 +140,7 @@ def _resolve_pcs_segment_lengths( pcs_segment_counts: tuple[int, ...], rigid_segment_selector: tuple[bool, ...], ) -> tuple[Array, ...]: - physical_lengths = jnp.asarray(params.length, dtype=jnp.float64) + physical_lengths = jnp.asarray(params.link.length, dtype=jnp.float64) pneumatic_physical_indices = [ i for i, is_rigid in enumerate(rigid_segment_selector) if not is_rigid ] @@ -202,18 +202,20 @@ def _expand_isupport_layout( params.validate() structure = _normalize_structure_for_params(params, structure) - num_physical_segments = int(params.length.shape[0]) - physical_lengths = jnp.asarray(params.length, dtype=jnp.float64) + num_physical_segments = int(params.link.length.shape[0]) + physical_lengths = jnp.asarray(params.link.length, dtype=jnp.float64) pcs_segment_lengths = _resolve_pcs_segment_lengths( params, structure.pcs_segment_counts, structure.rigid_segment_selector ) - radius = jnp.asarray(params.radius, dtype=jnp.float64) - density = jnp.asarray(params.density, dtype=jnp.float64) - young_modulus = jnp.asarray(params.young_modulus, dtype=jnp.float64) - shear_modulus = jnp.asarray(params.shear_modulus, dtype=jnp.float64) - reference_strain = jnp.asarray(params.reference_strain, dtype=jnp.float64).reshape( - num_physical_segments, 6 + radius = jnp.asarray( + params.link.cross_section.coefficients[:, 0], dtype=jnp.float64 ) + density = jnp.asarray(params.link.density, dtype=jnp.float64) + stiffness = jnp.asarray(params.link.stiffness, dtype=jnp.float64) + damping = jnp.asarray(params.link.damping, dtype=jnp.float64) + reference_strain = jnp.asarray( + params.link.reference_strain, dtype=jnp.float64 + ).reshape(num_physical_segments, 6) straight_reference_strain = _straight_reference_strain(reference_strain.dtype) for physical_index, is_rigid in enumerate(structure.rigid_segment_selector): if is_rigid and not bool( @@ -224,33 +226,6 @@ def _expand_isupport_layout( "[0, 0, 0, 1, 0, 0]; " f"physical segment {physical_index} does not." ) - has_damping_matrix = params.damping_matrix is not None - if has_damping_matrix: - damping_matrix = jnp.asarray(params.damping_matrix, dtype=jnp.float64) - damping_blocks = jnp.stack( - [ - damping_matrix[6 * i : 6 * (i + 1), 6 * i : 6 * (i + 1)] - for i in range(num_physical_segments) - ] - ) - if not bool(jnp.allclose(damping_matrix, blk_diag(damping_blocks))): - raise ValueError( - "ISupport damping_matrix must be block diagonal by physical segment " - "before expansion into PCS segments." - ) - else: - damping_matrix = None - damping_blocks = None - material_damping_coefficient = jnp.asarray( - params.material_damping_coefficient, dtype=jnp.float64 - ) - if material_damping_coefficient.ndim == 0: - material_damping_coefficient = jnp.full( - (num_physical_segments,), - material_damping_coefficient, - dtype=jnp.float64, - ) - r_chamber_in = jnp.asarray(params.chamber_inner_radius, dtype=jnp.float64) r_chamber_out = jnp.asarray(params.chamber_outer_radius, dtype=jnp.float64) d_chamber = jnp.asarray(params.chamber_distance, dtype=jnp.float64) @@ -261,11 +236,9 @@ def _expand_isupport_layout( expanded_lengths: list[float] = [] expanded_radius: list[Array] = [] expanded_density: list[Array] = [] - expanded_young_modulus: list[Array] = [] - expanded_shear_modulus: list[Array] = [] expanded_reference_strain: list[Array] = [] + expanded_stiffness: list[Array] = [] expanded_damping_blocks: list[Array] = [] - expanded_material_damping_coefficient: list[Array] = [] expanded_chamber_inner_radius: list[Array] = [] expanded_chamber_outer_radius: list[Array] = [] expanded_chamber_distance: list[Array] = [] @@ -284,8 +257,6 @@ def append_segment( expanded_lengths.append(jnp.asarray(length, dtype=jnp.float64)) expanded_radius.append(radius[source_index]) expanded_density.append(density[source_index]) - expanded_young_modulus.append(young_modulus[source_index]) - expanded_shear_modulus.append(shear_modulus[source_index]) chamber_index = 0 if pneumatic_index is None else pneumatic_index expanded_chamber_inner_radius.append(r_chamber_in[chamber_index]) expanded_chamber_outer_radius.append(r_chamber_out[chamber_index]) @@ -294,25 +265,15 @@ def append_segment( if is_rigid: expanded_reference_strain.append(reference_strain[source_index]) - if has_damping_matrix: - expanded_damping_blocks.append(damping_blocks[source_index]) - else: - expanded_material_damping_coefficient.append( - material_damping_coefficient[source_index] - ) + expanded_stiffness.append(stiffness[source_index]) + expanded_damping_blocks.append(damping[source_index]) default_strain_selector.append(jnp.zeros((6,), dtype=bool)) pcs_segment_to_pneumatic_segment.append(_RIGID_CONNECTOR_PARENT) else: length_scale = length / physical_lengths[source_index] expanded_reference_strain.append(reference_strain[source_index]) - if has_damping_matrix: - expanded_damping_blocks.append( - damping_blocks[source_index] * length_scale - ) - else: - expanded_material_damping_coefficient.append( - material_damping_coefficient[source_index] - ) + expanded_stiffness.append(stiffness[source_index] * length_scale) + expanded_damping_blocks.append(damping[source_index] * length_scale) default_strain_selector.append(jnp.ones((6,), dtype=bool)) pcs_segment_to_pneumatic_segment.append(pneumatic_index) pcs_segment_is_rigid.append(is_rigid) @@ -352,28 +313,23 @@ def append_segment( strain_selector ) - expanded_kwargs: dict[str, Array] = {} - if has_damping_matrix: - expanded_kwargs["damping_matrix"] = blk_diag(jnp.stack(expanded_damping_blocks)) - else: - expanded_kwargs["material_damping_coefficient"] = jnp.stack( - expanded_material_damping_coefficient - ) - expanded_params = ISupportParams( base_pose=params.base_pose, - length=jnp.stack(expanded_lengths), - radius=jnp.stack(expanded_radius), - density=jnp.stack(expanded_density), gravity=params.gravity, - young_modulus=jnp.stack(expanded_young_modulus), - shear_modulus=jnp.stack(expanded_shear_modulus), - reference_strain=jnp.concatenate(expanded_reference_strain), + link=ContinuumLinkParams( + length=jnp.stack(expanded_lengths), + density=jnp.stack(expanded_density), + reference_strain=jnp.stack(expanded_reference_strain), + cross_section=CrossSectionParams( + coefficients=jnp.stack(expanded_radius)[:, None] + ), + stiffness=jnp.stack(expanded_stiffness), + damping=jnp.stack(expanded_damping_blocks), + ), chamber_inner_radius=jnp.stack(expanded_chamber_inner_radius), chamber_outer_radius=jnp.stack(expanded_chamber_outer_radius), chamber_distance=jnp.stack(expanded_chamber_distance), chamber_azimuth_angles=jnp.stack(expanded_chamber_azimuth_angles), - **expanded_kwargs, ) pcs_structure = PCSStructure( num_gauss_points=structure.num_gauss_points, @@ -466,15 +422,19 @@ def __init__( structure: ISupportStructure | None = None, **kwargs: Any, ): - """ - Initialize the ISupport class + """Initialize an I-SUPPORT model from typed parameters. Args: params: Dynamic I-SUPPORT parameters. structure: Static I-SUPPORT layout. If omitted, physical segments alternate rigid and pneumatic from index zero, and each pneumatic segment is represented by one PCS segment. - **kwargs: Additional keyword arguments. + **kwargs: Additional keyword arguments forwarded to :class:`PCS`. + + Raises: + TypeError: If ``params`` or ``structure`` has the wrong type. + ValueError: If the physical layout, chamber geometry, expanded PCS + parameters, or an inherited constructor option is invalid. """ if not isinstance(params, ISupportParams): raise TypeError("params must be an ISupportParams instance.") @@ -611,7 +571,24 @@ def _current_body_params(self) -> ISupportParams: return getattr(self, "pcs_params", self.params) def with_params(self, params: ISupportParams) -> "ISupport": - """Return an updated copy with a full pneumatic-segment parameter object.""" + """Return a copy using complete physical-segment parameters. + + The physical I-SUPPORT layout is expanded back into PCS segments, and + inherited geometry, dynamics, and actuation caches are refreshed. + + Args: + params: Complete replacement :class:`ISupportParams` compatible + with the model's static physical layout and chamber count. + + Returns: + A new I-SUPPORT model containing the replacements. The original + model is unchanged. + + Raises: + TypeError: If ``params`` is not an :class:`ISupportParams`. + ValueError: If validation fails or the replacement changes a static + layout choice such as the number of chambers. + """ if not isinstance(params, ISupportParams): raise TypeError("params must be an ISupportParams instance.") params = _with_default_chamber_azimuth_angles(params) @@ -662,7 +639,22 @@ def with_params(self, params: ISupportParams) -> "ISupport": return updated_self._with_refreshed_precomputed_matrices() def update_params(self, **updates: Array) -> "ISupport": - """Return an updated copy with selected pneumatic parameter fields replaced.""" + """Return a copy with selected physical parameter fields replaced. + + Args: + **updates: Fields of :class:`ISupportParams` to replace, including + inherited PCS fields and pneumatic chamber fields. + + Returns: + A new validated I-SUPPORT model with refreshed PCS and actuator + caches. + + Raises: + KeyError: If an update names an unknown parameter field. + TypeError: If a replacement has an invalid type. + ValueError: If the resulting parameters are invalid or incompatible + with the static physical layout. + """ return self.with_params(self.params.replace(**updates)) @eqx.filter_jit @@ -831,6 +823,21 @@ def actuator_visual_layers( *, actuator_inputs: Array | None = None, ) -> tuple: - """Keep equivalent chamber-center paths out of generic renderers.""" + """Return generic actuator layers for rendering. + + I-SUPPORT's chamber-center paths represent equivalent pressure-volume + coordinates rather than physical tendons, so generic renderers should + not draw them. + + Args: + q: Generalized coordinates accepted by the common rendering API. + s_points: Backbone sampling coordinates accepted by the common + rendering API. + actuator_inputs: Optional actuator inputs accepted by the common + rendering API. + + Returns: + An empty tuple. + """ del q, s_points, actuator_inputs return () diff --git a/src/soromox/systems/pcs/params.py b/src/soromox/systems/pcs/params.py index 27b8a8be..09832308 100644 --- a/src/soromox/systems/pcs/params.py +++ b/src/soromox/systems/pcs/params.py @@ -6,10 +6,10 @@ from typing import ClassVar -import equinox as eqx import jax.numpy as jnp from jax import Array +from soromox.systems.components import ContinuumLinkParams from soromox.systems.params import ( BaseContinuumSoftRobotParams, validate_planar_base_pose, @@ -22,71 +22,40 @@ def _require_shape(name: str, value: Array, expected_shape: tuple[int, ...]) -> raise ValueError(f"{name} must have shape {expected_shape}, got {value.shape}.") -def _validate_damping_input( - params: "PCSParams | PlanarPCSParams", - *, - strain_dim: int, - n_segments: int, -) -> None: - has_damping_matrix = params.damping_matrix is not None - has_material_damping = params.material_damping_coefficient is not None - if has_damping_matrix == has_material_damping: - raise ValueError( - "Exactly one of damping_matrix or material_damping_coefficient " - "must be provided." - ) - - if has_damping_matrix: - damping_matrix = jnp.asarray(params.damping_matrix) - _require_shape( - "damping_matrix", - damping_matrix, - (strain_dim * n_segments, strain_dim * n_segments), - ) - return - - material_damping_coefficient = jnp.asarray(params.material_damping_coefficient) - if material_damping_coefficient.ndim == 0: - return - _require_shape( - "material_damping_coefficient", - material_damping_coefficient, - (n_segments,), - ) - - def _validate_continuum_base( params: BaseContinuumSoftRobotParams, *, strain_dim: int, gravity_dim: int, ) -> int: - if len(params.length.shape) != 1: - raise ValueError("length must be one-dimensional with shape (num_segments,).") - n_segments = params.length.shape[0] - if n_segments < 1: - raise ValueError(f"num_segments must be at least 1, got {n_segments}.") - _require_shape("density", params.density, (n_segments,)) + params.link.validate() + n_segments = params.link.length.shape[0] _require_shape("gravity", params.gravity, (gravity_dim,)) - expected_reference_size = strain_dim * n_segments - if params.reference_strain.size != expected_reference_size: + expected_reference_shape = (n_segments, strain_dim) + if params.link.reference_strain.shape != expected_reference_shape: raise ValueError( - "reference_strain must contain " - f"{expected_reference_size} entries, got {params.reference_strain.size}." + f"reference_strain must have shape {expected_reference_shape}, " + f"got {params.link.reference_strain.shape}." ) + expected_matrix_shape = (n_segments, strain_dim, strain_dim) + _require_shape("stiffness", params.link.stiffness, expected_matrix_shape) + _require_shape("damping", params.link.damping, expected_matrix_shape) + _require_shape( + "cross_section.coefficients", + params.link.cross_section.coefficients, + (n_segments, 1), + ) return n_segments class PCSParams(BaseContinuumSoftRobotParams): """Dynamic parameters for the spatial PCS model. - ``length``, ``radius``, material parameters, and density use a leading - segment axis. Damping can be supplied either as the preferred - ``material_damping_coefficient`` or as a full flattened strain - ``damping_matrix``. ``material_damping_coefficient`` is a viscosity-like - modulus in Pa*s (N*s/m^2); it may be scalar or have one value per segment. - The assembled matrix includes geometry and length factors, so its entries - have generalized-coordinate-dependent units rather than a single Pa*s unit. + Per-link runtime values live in the shared ``link`` component. Its + ``stiffness`` and ``damping`` fields are canonical generalized matrices with + shape ``(num_links, 6, 6)``. Isotropic Young, shear, and material-damping + values are separate construction or optimization variables and can be mapped + into these matrices explicitly. ``base_pose`` is the scalar-first quaternion SE(3) base pose vector ``[qw, qx, qy, qz, x, y, z]`` used to initialize the base transform. The quaternion is normalized before use and must have nonzero finite norm. @@ -96,57 +65,61 @@ class PCSParams(BaseContinuumSoftRobotParams): is_planar: ClassVar[bool] = False - radius: Array - young_modulus: Array - shear_modulus: Array - material_damping_coefficient: Array | None = eqx.field(default=None, kw_only=True) - damping_matrix: Array | None = eqx.field(default=None, kw_only=True) + link: ContinuumLinkParams def validate(self) -> None: - n_segments = _validate_continuum_base(self, strain_dim=6, gravity_dim=3) - _require_shape("radius", self.radius, (n_segments,)) - _require_shape("young_modulus", self.young_modulus, (n_segments,)) - _require_shape("shear_modulus", self.shear_modulus, (n_segments,)) - _validate_damping_input(self, strain_dim=6, n_segments=n_segments) + """Validate spatial PCS parameter shapes and values. + + Returns: + ``None`` after successful validation. + + Raises: + ValueError: If link fields are invalid; reference strain, canonical + matrices, cross-section coefficients, or gravity have the wrong + shape; or the spatial base pose is invalid. + """ + _validate_continuum_base(self, strain_dim=6, gravity_dim=3) validate_quaternion_base_pose("base_pose", self.base_pose, (7,)) class PlanarPCSParams(BaseContinuumSoftRobotParams): """Dynamic parameters for the planar PCS model. - The leading axis of per-segment fields indexes planar constant-strain - segments. ``base_pose`` stores the planar pose ``[theta, x, y]`` with shape + The leading axis of fields in ``link`` indexes planar constant-strain + segments. Its canonical generalized stiffness and damping matrices have + shape ``(num_links, 3, 3)``. ``base_pose`` stores the planar pose + ``[theta, x, y]`` with shape ``(3,)``. ``theta`` is a right-handed angle in radians about the out-of-plane z-axis, and ``x``/``y`` are direct translations in the parent - frame. ``material_damping_coefficient`` is a viscosity-like modulus in Pa*s - (N*s/m^2); it may be scalar or have one value per segment. The assembled - matrix includes geometry and length factors and therefore has - generalized-coordinate-dependent entry units. + frame. Isotropic material parameters are mapped explicitly into canonical + link matrices rather than duplicated in this runtime parameter tree. Omitting ``base_pose`` and ``gravity`` selects upright planar mounting and negative-y Earth gravity. """ is_planar: ClassVar[bool] = True - radius: Array - young_modulus: Array - shear_modulus: Array - material_damping_coefficient: Array | None = eqx.field(default=None, kw_only=True) - damping_matrix: Array | None = eqx.field(default=None, kw_only=True) + link: ContinuumLinkParams def validate(self) -> None: - n_segments = _validate_continuum_base(self, strain_dim=3, gravity_dim=2) - _require_shape("radius", self.radius, (n_segments,)) - _require_shape("young_modulus", self.young_modulus, (n_segments,)) - _require_shape("shear_modulus", self.shear_modulus, (n_segments,)) - _validate_damping_input(self, strain_dim=3, n_segments=n_segments) + """Validate planar PCS parameter shapes and values. + + Returns: + ``None`` after successful validation. + + Raises: + ValueError: If link fields are invalid; reference strain, canonical + matrices, cross-section coefficients, gravity, or the planar + base pose has the wrong shape or contains invalid values. + """ + _validate_continuum_base(self, strain_dim=3, gravity_dim=2) validate_planar_base_pose("base_pose", self.base_pose) class ISupportParams(PCSParams): """Dynamic parameters for the I-SUPPORT spatial pneumatic PCS model. - The leading axis of the standard PCS body fields indexes every physical + The leading axis of the shared PCS ``link`` fields indexes every physical rigid or pneumatic segment in robot order. ``ISupport`` expands these fields into the internal PCS layout using ``ISupportStructure``. Chamber fields index only pneumatic segments, in their order of appearance. Chamber @@ -157,11 +130,8 @@ class ISupportParams(PCSParams): length into its pressure-conjugate equivalent-volume coordinate. It contains one value per pneumatic segment and is shared by all chambers in that segment. When omitted, it is derived as - ``pi * (chamber_outer_radius**2 - chamber_inner_radius**2)``. Damping can be - supplied as ``material_damping_coefficient`` or as a full ``damping_matrix``. - A custom ``damping_matrix`` is expressed in flattened pneumatic-segment - strain coordinates and must be block diagonal by pneumatic segment when the - model is constructed. + ``pi * (chamber_outer_radius**2 - chamber_inner_radius**2)``. Stiffness and + damping are supplied as canonical per-link generalized matrices. ``pcs_segment_lengths`` optionally stores flattened PCS segment lengths in the order defined by ``ISupportStructure.pcs_segment_counts``. If omitted, @@ -180,6 +150,17 @@ class ISupportParams(PCSParams): chamber_effective_pressure_area: Array | None = None def validate(self) -> None: + """Validate PCS and pneumatic chamber parameters. + + Returns: + ``None`` after successful validation. + + Raises: + ValueError: If inherited PCS fields are invalid; chamber arrays + disagree in shape; radii, effective areas, or segment lengths + are invalid; or chamber azimuths are non-finite or not uniformly + distributed. + """ super().validate() chamber_inner_radius = jnp.asarray(self.chamber_inner_radius) if chamber_inner_radius.ndim != 1: diff --git a/src/soromox/systems/pcs/pcs.py b/src/soromox/systems/pcs/pcs.py index 45888fe7..c76a6977 100644 --- a/src/soromox/systems/pcs/pcs.py +++ b/src/soromox/systems/pcs/pcs.py @@ -1,7 +1,7 @@ __all__ = ["PCS"] -from typing import Any +from typing import Any, Self import equinox as eqx from jax import Array, lax, vmap @@ -12,9 +12,16 @@ BaseThreadlikeRoutingParams, ThreadlikeRouting, ) +from soromox.systems.components import ( + ContinuumLinkParams, + CrossSectionGeometry, + CrossSectionParams, + IsotropicMaterialParams, + LinkSpec, +) from soromox.systems.pcs.params import PCSParams from soromox.systems.pcs.structures import PCSStructure -from soromox.systems.soft_robot import CrossSectionGeometry, SoftRobot +from soromox.systems.soft_robot import SoftRobot from soromox.utils.array_math import blk_diag from soromox.utils.basic import ( compute_strain_basis, @@ -89,8 +96,6 @@ class PCS(SoftRobot): L_cum: Array # Cumulative length of the segments r: Array # Radius of the segments rho: Array - E: Array # Young's modulus of the segments - G: Array # Shear modulus of the segments num_segments: int = eqx.field(static=True) num_gauss_points: int = eqx.field(static=True) @@ -110,6 +115,166 @@ class PCS(SoftRobot): K_active: Array # Cached active-coordinate stiffness matrix D_full: Array # Cached full damping matrix D_active: Array # Cached active-coordinate damping matrix + young_stiffness_operator: Array + shear_stiffness_operator: Array + material_damping_operator: Array + + @staticmethod + def params_from_links( + links: list[LinkSpec] | tuple[LinkSpec, ...], + *, + gravity: Array | None = None, + base_pose: Array | None = None, + ) -> PCSParams: + """Build spatial PCS parameters from shared link specifications. + + Args: + links: Non-empty sequence of constant, circular link + specifications. Each reference strain must have shape ``(6,)``. + gravity: Optional world-frame gravity vector with shape ``(3,)``. + base_pose: Optional base translation and unit quaternion with shape + ``(7,)``. + + Returns: + Validated :class:`PCSParams` with canonical per-link stiffness and + damping arrays of shape ``(num_links, 6, 6)``. + + Raises: + ValueError: If no links are supplied; a link is not constant and + circular; a reference strain has the wrong shape; a material + source is incomplete; or an explicit matrix is not ``(6, 6)``. + """ + if not links: + raise ValueError("PCS requires at least one link.") + specs = tuple(links) + for index, spec in enumerate(specs): + if spec.cross_section_geometry != CrossSectionGeometry.CIRCULAR: + raise ValueError(f"PCS link {index} must have a circular section.") + if spec.cross_section_profile_types != ("constant",): + raise ValueError(f"PCS link {index} must have a constant radius.") + if jnp.asarray(spec.reference_strain).shape != (6,): + raise ValueError( + f"PCS link {index} reference_strain must have shape (6,)." + ) + + length = jnp.asarray([spec.length for spec in specs]) + density = jnp.asarray([spec.density for spec in specs]) + radius = jnp.asarray([spec.cross_section_coefficients[0] for spec in specs]) + reference_strain = jnp.asarray( + [spec.reference_strain for spec in specs], dtype=float + ) + area = jnp.pi * radius**2 + transverse = jnp.pi * radius**4 / 4.0 + polar = 2.0 * transverse + young_operator = length[:, None, None] * vmap(jnp.diag)( + jnp.stack( + [ + jnp.zeros_like(length), + transverse, + transverse, + area, + jnp.zeros_like(length), + jnp.zeros_like(length), + ], + axis=1, + ) + ) + shear_operator = length[:, None, None] * vmap(jnp.diag)( + jnp.stack( + [ + polar, + jnp.zeros_like(length), + jnp.zeros_like(length), + jnp.zeros_like(length), + area, + area, + ], + axis=1, + ) + ) + damping_operator = length[:, None, None] * vmap(jnp.diag)( + jnp.stack( + [polar, 3.0 * transverse, 3.0 * transverse, 3.0 * area, area, area], + axis=1, + ) + ) + + stiffness_items = [] + damping_items = [] + for index, spec in enumerate(specs): + if spec.stiffness is not None: + stiffness = jnp.asarray(spec.stiffness) + if stiffness.shape != (6, 6): + raise ValueError( + f"PCS link {index} stiffness must have shape (6, 6)." + ) + else: + stiffness = ( + spec.young_modulus * young_operator[index] + + spec.shear_modulus * shear_operator[index] + ) + if spec.damping is not None: + damping = jnp.asarray(spec.damping) + if damping.shape != (6, 6): + raise ValueError( + f"PCS link {index} damping must have shape (6, 6)." + ) + else: + damping = spec.material_damping_coefficient * damping_operator[index] + stiffness_items.append(stiffness) + damping_items.append(damping) + + return PCSParams( + base_pose=base_pose, + gravity=gravity, + link=ContinuumLinkParams( + length=length, + density=density, + reference_strain=reference_strain, + cross_section=CrossSectionParams(coefficients=radius[:, None]), + stiffness=jnp.stack(stiffness_items), + damping=jnp.stack(damping_items), + ), + ) + + @classmethod + def from_links( + cls, + links: list[LinkSpec] | tuple[LinkSpec, ...], + *, + structure: PCSStructure | None = None, + gravity: Array | None = None, + base_pose: Array | None = None, + **kwargs: Any, + ) -> "PCS": + """Construct a spatial PCS model from shared link specifications. + + Args: + links: Non-empty sequence of constant, circular link + specifications. + structure: Optional static PCS quadrature, strain-selection, and + rotational-scaling configuration. Defaults to + :class:`PCSStructure`. + gravity: Optional world-frame gravity vector with shape ``(3,)``. + base_pose: Optional base translation and unit quaternion with shape + ``(7,)``. + **kwargs: Additional keyword arguments forwarded to the PCS + constructor, such as actuators or passive elements. + + Returns: + A fully initialized spatial PCS model. + + Raises: + TypeError: If a structure or constructor argument has an invalid + type. + ValueError: If a link specification, static structure, parameter + array, or attached component is invalid. + """ + return cls( + params=cls.params_from_links(links, gravity=gravity, base_pose=base_pose), + structure=structure, + **kwargs, + ) def __init__( self, @@ -119,7 +284,25 @@ def __init__( passive_elements: PassiveElement | tuple[PassiveElement, ...] | None = (), **kwargs: Any, ): - """Initialize the PCS class from typed dynamic parameters.""" + """Initialize a spatial PCS model from typed parameters. + + Args: + params: Canonical dynamic link, gravity, and base-pose parameters. + structure: Optional static quadrature, active-strain, and + rotational-scaling configuration. + actuators: Optional actuator or tuple of actuators attached to the + model. + passive_elements: Optional passive element or tuple of passive + elements. Pass ``None`` or an empty tuple to disable them. + **kwargs: Additional keyword arguments forwarded to + :class:`BaseContinuumSoftRobot`. + + Raises: + TypeError: If ``params``, a structure field, or an attached + component has an invalid type. + ValueError: If parameter shapes or values, the structure, or an + attached component are invalid. + """ if not isinstance(params, PCSParams): raise TypeError("params must be a PCSParams instance.") params.validate() @@ -132,7 +315,7 @@ def __init__( ) # Number of segments - num_segments = int(params.length.shape[0]) + num_segments = int(params.link.length.shape[0]) if num_segments < 1: raise ValueError(f"num_segments must be at least 1, got {num_segments}") self.num_segments = num_segments @@ -189,7 +372,7 @@ def __init__( self.num_active_strains = jnp.sum(strain_selector) self.num_dofs = int(self.num_active_strains.item()) - reference_strain = jnp.asarray(params.reference_strain, dtype=jnp.float64) + reference_strain = jnp.asarray(params.link.reference_strain, dtype=jnp.float64) if reference_strain.size != num_strains: raise ValueError( "reference_strain must have " @@ -203,12 +386,20 @@ def __init__( @property def is_planar(self) -> bool: - """PCS is a spatial (3D) model.""" + """Return whether the system is planar. + + Returns: + Always ``False`` because spatial PCS uses SE(3) kinematics. + """ return False @property def segment_length(self) -> Array: - """Per-segment backbone lengths.""" + """Return the per-segment backbone lengths. + + Returns: + Array with shape ``(num_segments,)``. + """ return jnp.asarray(self.L) def _strain_basis_scaling_vector(self) -> Array: @@ -232,7 +423,21 @@ def _scaled_strain_basis(self, B_xi: Array) -> Array: return B_xi def cross_section_geometry(self, q: Array, s: Array) -> tuple[Array, Array]: - """Return the assumed solid circular cross-section and segment radius.""" + """Evaluate the circular cross-section at a backbone coordinate. + + Args: + q: Generalized coordinates. PCS cross-section geometry is + configuration-independent, but ``q`` is accepted for the common + continuum-robot interface. + s: Scalar global backbone coordinate. + + Returns: + A tuple containing the circular :class:`CrossSectionGeometry` tag + and a one-element array holding the selected segment radius. + + Raises: + ValueError: If ``s`` cannot be assigned to a segment. + """ segment_idx, _ = self.classify_segment(s) radius = jnp.asarray(self.r)[segment_idx] tag = jnp.asarray(CrossSectionGeometry.CIRCULAR, dtype=jnp.int32) @@ -258,7 +463,7 @@ def _set_params(self, params: PCSParams) -> None: ) # Add zeros for the orientation angles # Lengths of the segments - L = params.length + L = params.link.length L = jnp.asarray(L, dtype=jnp.float64) if L.shape != (self.num_segments,): raise ValueError( @@ -270,16 +475,11 @@ def _set_params(self, params: PCSParams) -> None: self.L_cum = L_cum # Radius of the segments - r = params.radius - r = jnp.asarray(r, dtype=jnp.float64) - if r.shape != (self.num_segments,): - raise ValueError( - f"radius must have shape ({self.num_segments},), got {r.shape}" - ) + r = jnp.asarray(params.link.cross_section.coefficients[:, 0], dtype=jnp.float64) self.r = r # Densities of the segments - rho = params.density + rho = params.link.density rho = jnp.asarray(rho, dtype=jnp.float64) if rho.shape != (self.num_segments,): raise ValueError( @@ -287,82 +487,9 @@ def _set_params(self, params: PCSParams) -> None: ) self.rho = rho - # Elastic modulus of the segments - E = params.young_modulus - E = jnp.asarray(E, dtype=jnp.float64) - if E.shape != (self.num_segments,): - raise ValueError( - f"young_modulus must have shape ({self.num_segments},), got {E.shape}" - ) - self.E = E - - # Shear modulus of the segments - G = params.shear_modulus - G = jnp.asarray(G, dtype=jnp.float64) - if G.shape != (self.num_segments,): - raise ValueError( - f"shear_modulus must have shape ({self.num_segments},), got {G.shape}" - ) - self.G = G - - def _explicit_damping_full_matrix(self, params: PCSParams) -> Array: - """Return the custom full damping matrix supplied in params.""" - if params.damping_matrix is None: - raise ValueError("damping_matrix is not set.") - expected_D_shape = (self.num_strains, self.num_strains) - D = jnp.asarray(params.damping_matrix, dtype=jnp.float64) - if D.shape != expected_D_shape: - raise ValueError( - f"damping_matrix must have shape {expected_D_shape}, got {D.shape}" - ) - return D - - def _material_damping_coefficients(self) -> Array: - """Return per-segment material damping coefficients in Pa*s (N*s/m^2).""" - params = self._current_body_params() - if params.material_damping_coefficient is None: - raise ValueError("material_damping_coefficient is not set.") - coefficient = jnp.asarray( - params.material_damping_coefficient, dtype=jnp.float64 - ) - if coefficient.ndim == 0: - return jnp.full((self.num_segments,), coefficient, dtype=jnp.float64) - if coefficient.shape != (self.num_segments,): - raise ValueError( - "material_damping_coefficient must be a scalar or have shape " - f"({self.num_segments},), got {coefficient.shape}." - ) - return coefficient - - def _compute_material_damping_full_matrix(self) -> Array: - """Compute damping using the solid-circle area and second moments.""" - coefficients = self._material_damping_coefficients() - - def damping_block(i: Array) -> Array: - I_i = self._local_second_moment_of_area(i) - A_i = self._local_cross_sectional_area(i) - damping_diag = jnp.stack( - [ - I_i[0], - 3.0 * I_i[1], - 3.0 * I_i[2], - 3.0 * A_i, - A_i, - A_i, - ], - axis=0, - ) - return self.L[i] * coefficients[i] * jnp.diag(damping_diag) - - damping_blocks = vmap(damping_block)(jnp.arange(self.num_segments)) - return blk_diag(damping_blocks) - def _compute_damping_full_matrix(self) -> Array: - """Compute the current full damping matrix.""" - params = self._current_body_params() - if params.material_damping_coefficient is not None: - return self._compute_material_damping_full_matrix() - return self._explicit_damping_full_matrix(params) + """Assemble canonical per-link generalized damping blocks.""" + return blk_diag(self._current_body_params().link.damping) def _current_body_params(self) -> PCSParams: """Return the PCS body params, including for typed actuated wrappers.""" @@ -381,23 +508,26 @@ def _with_pcs_params( if not isinstance(params, PCSParams): raise TypeError("params must be a PCSParams instance.") params.validate() - if params.length.shape != current_params.length.shape: + if params.link.length.shape != current_params.link.length.shape: raise ValueError( "length shape changes the model structure; construct a new PCS." ) - if params.reference_strain.shape != current_params.reference_strain.shape: + if ( + params.link.reference_strain.shape + != current_params.link.reference_strain.shape + ): raise ValueError( "reference_strain shape changes the model structure; construct a new PCS." ) base_pose = jnp.asarray(params.base_pose, dtype=jnp.float64) gravity = jnp.asarray(params.gravity, dtype=jnp.float64) - segment_lengths = jnp.asarray(params.length, dtype=jnp.float64) - radius = jnp.asarray(params.radius, dtype=jnp.float64) - density = jnp.asarray(params.density, dtype=jnp.float64) - young_modulus = jnp.asarray(params.young_modulus, dtype=jnp.float64) - shear_modulus = jnp.asarray(params.shear_modulus, dtype=jnp.float64) - reference_strain = jnp.asarray(params.reference_strain, dtype=jnp.float64) + segment_lengths = jnp.asarray(params.link.length, dtype=jnp.float64) + radius = jnp.asarray( + params.link.cross_section.coefficients[:, 0], dtype=jnp.float64 + ) + density = jnp.asarray(params.link.density, dtype=jnp.float64) + reference_strain = jnp.asarray(params.link.reference_strain, dtype=jnp.float64) updated_self = eqx.tree_at( lambda m: ( @@ -409,8 +539,6 @@ def _with_pcs_params( m.L_cum, m.r, m.rho, - m.E, - m.G, m.xi_ref, ), self, @@ -427,36 +555,132 @@ def _with_pcs_params( ), radius, density, - young_modulus, - shear_modulus, reference_strain.reshape(self.num_strains), ), ) return updated_self._with_refreshed_precomputed_matrices() def with_params(self, params: PCSParams) -> "PCS": - """Return an updated copy with a full typed parameter object.""" + """Return a model copy using a complete parameter PyTree. + + Args: + params: Complete replacement parameters with the same number of + segments and reference-strain layout as this model. + + Returns: + A new PCS model with refreshed geometry, mass, material-response, + stiffness, and damping caches. The original model is unchanged. + + Raises: + TypeError: If ``params`` is not a :class:`PCSParams`. + ValueError: If the replacement is invalid or changes the static + segment or strain layout. + """ return self._with_pcs_params(params) - def update_params(self, **updates: Array) -> "PCS": - """Return an updated copy with selected typed parameter fields replaced.""" + def update_params(self, **updates: Any) -> "PCS": + """Return a copy with selected top-level parameter fields replaced. + + Args: + **updates: Fields of :class:`PCSParams` to replace, typically + ``link``, ``gravity``, or ``base_pose``. + + Returns: + A new validated PCS model containing the replacements. + + Raises: + TypeError: If an unknown field is supplied or a replacement has an + invalid type. + ValueError: If the result is invalid or changes the static segment + or strain layout. + """ if ( - "length" in updates - and jnp.asarray(updates["length"]).shape != self.params.length.shape + "link" in updates + and jnp.asarray(updates["link"].length).shape + != self.params.link.length.shape ): raise ValueError( "length shape changes the model structure; construct a new PCS." ) if ( - "reference_strain" in updates - and jnp.asarray(updates["reference_strain"]).shape - != self.params.reference_strain.shape + "link" in updates + and jnp.asarray(updates["link"].reference_strain).shape + != self.params.link.reference_strain.shape ): raise ValueError( "reference_strain shape changes the model structure; construct a new PCS." ) return self.with_params(self.params.replace(**updates)) + def update_link_params(self, **updates: Any) -> "PCS": + """Return a copy with selected continuum-link fields replaced. + + Args: + **updates: Fields of :class:`ContinuumLinkParams` to replace, such + as ``length``, ``density``, ``reference_strain``, + ``cross_section``, ``stiffness``, or ``damping``. + + Returns: + A new validated PCS model with refreshed dependent caches. + + Raises: + TypeError: If an unknown link field is supplied. + ValueError: If a replacement is invalid or changes the static + segment or strain layout. + """ + return self.with_params( + self.params.replace(link=self.params.link.replace(**updates)) + ) + + def link_matrices_from_material( + self, material: IsotropicMaterialParams + ) -> tuple[Array, Array]: + """Map isotropic material variables to canonical PCS link matrices. + + Args: + material: Scalar or per-segment Young's modulus, shear modulus, and + material damping coefficient. Scalars are broadcast over all + segments. + + Returns: + A tuple ``(stiffness, damping)`` whose arrays both have shape + ``(num_segments, 6, 6)``. + + Raises: + ValueError: If a material field is not scalar or does not have + shape ``(num_segments,)``. + """ + material = material.broadcast(self.num_segments) + stiffness = ( + material.young_modulus[:, None, None] * self.young_stiffness_operator + + material.shear_modulus[:, None, None] * self.shear_stiffness_operator + ) + damping = ( + material.material_damping_coefficient[:, None, None] + * self.material_damping_operator + ) + return stiffness, damping + + def with_isotropic_material(self, material: IsotropicMaterialParams) -> Self: + """Return a copy whose link matrices are built from isotropic material. + + The supplied material PyTree remains caller-owned and is not stored on + the model. + + Args: + material: Scalar or per-segment isotropic material variables. + + Returns: + A new PCS model containing the generated canonical link stiffness + and damping matrices. + + Raises: + ValueError: If a material field cannot be broadcast to one value + per segment. + """ + stiffness, damping = self.link_matrices_from_material(material) + return self.update_link_params(stiffness=stiffness, damping=damping) + def _precomputed_matrices(self) -> tuple[Array, Array, Array, Array, Array]: """Compute state-independent matrices cached by the model.""" M_segments = vmap(self._compute_local_mass_matrix)( @@ -469,8 +693,17 @@ def _precomputed_matrices(self) -> tuple[Array, Array, Array, Array, Array]: return M_segments, K_full, K_active, D_full, D_active def precompute(self) -> None: - """Refresh state-independent matrices cached by the model.""" + """Refresh state-independent matrices cached by the model. + + Returns: + ``None``. Material unit-response operators and the segment, full, + and active-coordinate matrices are replaced in place. + """ object.__setattr__(self, "B_xi", self._scaled_strain_basis(self.B_xi_unscaled)) + young_operator, shear_operator, damping_operator = self._material_operators() + object.__setattr__(self, "young_stiffness_operator", young_operator) + object.__setattr__(self, "shear_stiffness_operator", shear_operator) + object.__setattr__(self, "material_damping_operator", damping_operator) ( M_segments, K_full, @@ -488,6 +721,18 @@ def _with_refreshed_precomputed_matrices(self) -> "PCS": """Return a copy with cached state-independent matrices refreshed.""" B_xi = self._scaled_strain_basis(self.B_xi_unscaled) updated_self = eqx.tree_at(lambda m: m.B_xi, self, B_xi) + young_operator, shear_operator, damping_operator = ( + updated_self._material_operators() + ) + updated_self = eqx.tree_at( + lambda m: ( + m.young_stiffness_operator, + m.shear_stiffness_operator, + m.material_damping_operator, + ), + updated_self, + (young_operator, shear_operator, damping_operator), + ) ( M_segments, K_full, @@ -1316,8 +1561,17 @@ def skip_branch( def jacobian_and_arc_length_derivative_bodyframe( self, q: Array, s: Array ) -> tuple[Array, Array]: - """ - Compute the body-frame Jacobian and its arc-length derivative at ``s``. + """Compute a body-frame Jacobian and its arc-length derivative. + + Args: + q: Active generalized strains with shape + ``(num_active_strains,)``. + s: Scalar global backbone coordinate. + + Returns: + A tuple ``(J, J_s)`` containing the body-frame Jacobian and its + derivative with respect to global arc length. Both arrays have + shape ``(6, num_active_strains)``. """ xi = self.strain(q).reshape(self.num_segments, 6) segment_idx, s_local = self.classify_segment(s) @@ -1400,8 +1654,16 @@ def skip_branch( @eqx.filter_jit def jacobian_arc_length_derivative_bodyframe(self, q: Array, s: Array) -> Array: - """ - Compute the arc-length derivative of the body-frame Jacobian at ``s``. + """Compute the body-frame Jacobian derivative with respect to arc length. + + Args: + q: Active generalized strains with shape + ``(num_active_strains,)``. + s: Scalar global backbone coordinate. + + Returns: + Arc-length derivative with shape + ``(6, num_active_strains)``. """ _, Js = self.jacobian_and_arc_length_derivative_bodyframe(q, s) return Js @@ -1456,8 +1718,17 @@ def jacobian_inertialframe(self, q: Array, s: Array) -> Array: def jacobian_and_arc_length_derivative_inertialframe( self, q: Array, s: Array ) -> tuple[Array, Array]: - """ - Compute the inertial-frame Jacobian and its arc-length derivative at ``s``. + """Compute an inertial-frame Jacobian and its arc-length derivative. + + Args: + q: Active generalized strains with shape + ``(num_active_strains,)``. + s: Scalar global backbone coordinate. + + Returns: + A tuple ``(J, J_s)`` containing the inertial-frame Jacobian and its + derivative with respect to global arc length. Both arrays have + shape ``(6, num_active_strains)``. """ g_s, J_local, Js_local = ( self._jacobian_and_arc_length_derivative_bodyframe_with_pose(q, s) @@ -1475,8 +1746,16 @@ def jacobian_and_arc_length_derivative_inertialframe( @eqx.filter_jit def jacobian_arc_length_derivative_inertialframe(self, q: Array, s: Array) -> Array: - """ - Compute the arc-length derivative of the inertial-frame Jacobian at ``s``. + """Compute the inertial Jacobian derivative with respect to arc length. + + Args: + q: Active generalized strains with shape + ``(num_active_strains,)``. + s: Scalar global backbone coordinate. + + Returns: + Arc-length derivative with shape + ``(6, num_active_strains)``. """ g_s, J_local, Js_local = ( self._jacobian_and_arc_length_derivative_bodyframe_with_pose(q, s) @@ -1934,7 +2213,17 @@ def _jacobian_and_arc_length_derivative( @eqx.filter_jit def jacobian_batched(self, q: Array, s_ps: Array) -> Array: - """Compute inertial-frame Jacobians at multiple arc-length positions.""" + """Compute inertial-frame Jacobians at multiple arc-length positions. + + Args: + q: Active generalized strains with shape + ``(num_active_strains,)``. + s_ps: Backbone coordinates with shape ``(num_points,)``. + + Returns: + Inertial-frame Jacobians with shape + ``(num_points, 6, num_active_strains)``. + """ return self.jacobian_inertialframe_batched(q, s_ps) @eqx.filter_jit @@ -1948,7 +2237,19 @@ def _jacobian_and_time_derivative( def jacobian_and_time_derivative_batched( self, q: Array, qd: Array, s_ps: Array ) -> tuple[Array, Array]: - """Compute inertial-frame Jacobians and time derivatives at multiple arc-length positions.""" + """Compute batched inertial Jacobians and their time derivatives. + + Args: + q: Active generalized strains with shape + ``(num_active_strains,)``. + qd: Active generalized strain rates with shape + ``(num_active_strains,)``. + s_ps: Backbone coordinates with shape ``(num_points,)``. + + Returns: + A tuple ``(J, J_dot)`` whose arrays both have shape + ``(num_points, 6, num_active_strains)``. + """ return self.jacobian_and_time_derivative_inertialframe_batched(q, qd, s_ps) # ========================================== @@ -2285,47 +2586,29 @@ def _gravitational_force(self, q: Array) -> Array: return G - @eqx.filter_jit - def _local_stiffness_matrix(self, i: Array) -> Array: - """ - Compute local stiffness for a spatial rod longitudinally aligned with local x. - - The diagonal constitutive terms use the solid circular cross-section's - area and second moments. - - Args: - i (Array): index of the segment as array of shape () - - Returns: - S_i (Array): Local stiffness matrix of shape (6, 6) for the i-th segment. - """ - I_i = self._local_second_moment_of_area( - i - ) # Second moment of area as array of shape (3, ) - A_i = self._local_cross_sectional_area(i) # Cross-sectional area - - S_i = self.L[i] * jnp.diag( + def _material_operators(self) -> tuple[Array, Array, Array]: + """Return unit Young, shear, and material-damping link matrices.""" + area = jnp.pi * self.r**2 + transverse = jnp.pi * self.r**4 / 4.0 + polar = 2.0 * transverse + zeros = jnp.zeros_like(self.L) + young = self.L[:, None, None] * vmap(jnp.diag)( + jnp.stack([zeros, transverse, transverse, area, zeros, zeros], axis=1) + ) + shear = self.L[:, None, None] * vmap(jnp.diag)( + jnp.stack([polar, zeros, zeros, zeros, area, area], axis=1) + ) + damping = self.L[:, None, None] * vmap(jnp.diag)( jnp.stack( - [ - self.G[i] * I_i[0], # torsion X - self.E[i] * I_i[1], # bending Y - self.E[i] * I_i[2], # bending Z - A_i * self.E[i], # axial X - A_i * self.G[i], # shear Y - A_i * self.G[i], # shear Z - ], - axis=0, + [polar, 3.0 * transverse, 3.0 * transverse, 3.0 * area, area, area], + axis=1, ) ) - - return S_i + return young, shear, damping def _compute_stiffness_full_matrix(self) -> Array: - """Compute the uncached full stiffness matrix from current parameters.""" - # stiffness matrix of shape (num_segments, 6, 6) - S_sms = vmap(self._local_stiffness_matrix)(jnp.arange(self.num_segments)) - # we define the elastic matrix of shape (num_strains, num_strains) as K(xi) = K @ xi where K is equal to - return blk_diag(S_sms) + """Assemble canonical per-link generalized stiffness blocks.""" + return blk_diag(self._current_body_params().link.stiffness) @eqx.filter_jit def _stiffness(self, formulate_in_strain_space: bool = False) -> Array: diff --git a/src/soromox/systems/pcs/planar_pcs.py b/src/soromox/systems/pcs/planar_pcs.py index 55a26fc4..eaf0e1d9 100644 --- a/src/soromox/systems/pcs/planar_pcs.py +++ b/src/soromox/systems/pcs/planar_pcs.py @@ -1,7 +1,7 @@ __all__ = ["PlanarPCS"] -from typing import Any +from typing import Any, Self import equinox as eqx from jax import Array, lax, vmap @@ -12,9 +12,16 @@ BaseThreadlikeRoutingParams, ThreadlikeRouting, ) +from soromox.systems.components import ( + ContinuumLinkParams, + CrossSectionGeometry, + CrossSectionParams, + IsotropicMaterialParams, + LinkSpec, +) from soromox.systems.pcs.params import PlanarPCSParams from soromox.systems.pcs.structures import PlanarPCSStructure -from soromox.systems.soft_robot import CrossSectionGeometry, SoftRobot +from soromox.systems.soft_robot import SoftRobot from soromox.utils.array_math import blk_diag from soromox.utils.basic import ( compute_strain_basis, @@ -86,8 +93,6 @@ class PlanarPCS(SoftRobot): L_cum: Array # Cumulative length of the segments r: Array # Radius of the segments rho: Array - E: Array # Young's modulus of the segments - G: Array # Shear modulus of the segments num_segments: int = eqx.field(static=True) num_gauss_points: int = eqx.field(static=True) @@ -107,6 +112,142 @@ class PlanarPCS(SoftRobot): K_active: Array # Cached active-coordinate stiffness matrix D_full: Array # Cached full damping matrix D_active: Array # Cached active-coordinate damping matrix + young_stiffness_operator: Array + shear_stiffness_operator: Array + material_damping_operator: Array + + @staticmethod + def params_from_links( + links: list[LinkSpec] | tuple[LinkSpec, ...], + *, + gravity: Array | None = None, + base_pose: Array | None = None, + ) -> PlanarPCSParams: + """Build planar PCS parameters from shared link specifications. + + Args: + links: Non-empty sequence of constant, circular link + specifications. Each reference strain must have shape ``(3,)``. + gravity: Optional planar gravity vector with shape ``(2,)``. + base_pose: Optional planar base pose ``[theta, x, y]`` with shape + ``(3,)``. + + Returns: + Validated :class:`PlanarPCSParams` with canonical per-link + stiffness and damping arrays of shape ``(num_links, 3, 3)``. + + Raises: + ValueError: If no links are supplied; a link is not constant and + circular; a reference strain has the wrong shape; a material + source is incomplete; or an explicit matrix is not ``(3, 3)``. + """ + if not links: + raise ValueError("PlanarPCS requires at least one link.") + specs = tuple(links) + for index, spec in enumerate(specs): + if spec.cross_section_geometry != CrossSectionGeometry.CIRCULAR: + raise ValueError( + f"PlanarPCS link {index} must have a circular section." + ) + if spec.cross_section_profile_types != ("constant",): + raise ValueError(f"PlanarPCS link {index} must have a constant radius.") + if jnp.asarray(spec.reference_strain).shape != (3,): + raise ValueError( + f"PlanarPCS link {index} reference_strain must have shape (3,)." + ) + length = jnp.asarray([spec.length for spec in specs]) + density = jnp.asarray([spec.density for spec in specs]) + radius = jnp.asarray([spec.cross_section_coefficients[0] for spec in specs]) + reference_strain = jnp.asarray( + [spec.reference_strain for spec in specs], dtype=float + ) + area = jnp.pi * radius**2 + moment = jnp.pi * radius**4 / 4.0 + zeros = jnp.zeros_like(length) + young_operator = length[:, None, None] * vmap(jnp.diag)( + jnp.stack([moment, area, zeros], axis=1) + ) + shear_operator = length[:, None, None] * vmap(jnp.diag)( + jnp.stack([zeros, zeros, area], axis=1) + ) + damping_operator = length[:, None, None] * vmap(jnp.diag)( + jnp.stack([3.0 * moment, 3.0 * area, area], axis=1) + ) + stiffness_items = [] + damping_items = [] + for index, spec in enumerate(specs): + if spec.stiffness is not None: + stiffness = jnp.asarray(spec.stiffness) + if stiffness.shape != (3, 3): + raise ValueError( + f"PlanarPCS link {index} stiffness must have shape (3, 3)." + ) + else: + stiffness = ( + spec.young_modulus * young_operator[index] + + spec.shear_modulus * shear_operator[index] + ) + if spec.damping is not None: + damping = jnp.asarray(spec.damping) + if damping.shape != (3, 3): + raise ValueError( + f"PlanarPCS link {index} damping must have shape (3, 3)." + ) + else: + damping = spec.material_damping_coefficient * damping_operator[index] + stiffness_items.append(stiffness) + damping_items.append(damping) + return PlanarPCSParams( + base_pose=base_pose, + gravity=gravity, + link=ContinuumLinkParams( + length=length, + density=density, + reference_strain=reference_strain, + cross_section=CrossSectionParams(coefficients=radius[:, None]), + stiffness=jnp.stack(stiffness_items), + damping=jnp.stack(damping_items), + ), + ) + + @classmethod + def from_links( + cls, + links: list[LinkSpec] | tuple[LinkSpec, ...], + *, + structure: PlanarPCSStructure | None = None, + gravity: Array | None = None, + base_pose: Array | None = None, + **kwargs: Any, + ) -> "PlanarPCS": + """Construct a planar PCS model from shared link specifications. + + Args: + links: Non-empty sequence of constant, circular link + specifications. + structure: Optional static planar PCS quadrature, strain-selection, + and rotational-scaling configuration. Defaults to + :class:`PlanarPCSStructure`. + gravity: Optional planar gravity vector with shape ``(2,)``. + base_pose: Optional planar base pose ``[theta, x, y]`` with shape + ``(3,)``. + **kwargs: Additional keyword arguments forwarded to the PlanarPCS + constructor, such as actuators or passive elements. + + Returns: + A fully initialized planar PCS model. + + Raises: + TypeError: If a structure or constructor argument has an invalid + type. + ValueError: If a link specification, static structure, parameter + array, or attached component is invalid. + """ + return cls( + params=cls.params_from_links(links, gravity=gravity, base_pose=base_pose), + structure=structure, + **kwargs, + ) def __init__( self, @@ -116,7 +257,25 @@ def __init__( passive_elements: PassiveElement | tuple[PassiveElement, ...] | None = (), **kwargs: Any, ): - """Initialize the PlanarPCS class from typed dynamic parameters.""" + """Initialize a planar PCS model from typed parameters. + + Args: + params: Canonical dynamic link, gravity, and base-pose parameters. + structure: Optional static quadrature, active-strain, and + rotational-scaling configuration. + actuators: Optional actuator or tuple of actuators attached to the + model. + passive_elements: Optional passive element or tuple of passive + elements. Pass ``None`` or an empty tuple to disable them. + **kwargs: Additional keyword arguments forwarded to + :class:`BaseContinuumSoftRobot`. + + Raises: + TypeError: If ``params``, a structure field, or an attached + component has an invalid type. + ValueError: If parameter shapes or values, the structure, or an + attached component are invalid. + """ if not isinstance(params, PlanarPCSParams): raise TypeError("params must be a PlanarPCSParams instance.") params.validate() @@ -129,7 +288,7 @@ def __init__( ) # Number of segments - num_segments = int(params.length.shape[0]) + num_segments = int(params.link.length.shape[0]) if num_segments < 1: raise ValueError(f"num_segments must be at least 1, got {num_segments}") self.num_segments = num_segments @@ -186,7 +345,7 @@ def __init__( self.num_active_strains = jnp.sum(strain_selector) self.num_dofs = int(self.num_active_strains.item()) - reference_strain = jnp.asarray(params.reference_strain, dtype=jnp.float64) + reference_strain = jnp.asarray(params.link.reference_strain, dtype=jnp.float64) if reference_strain.size != self.num_strains: raise ValueError( "reference_strain must have " @@ -200,12 +359,20 @@ def __init__( @property def is_planar(self) -> bool: - """Planar PCS is a 2D model.""" + """Return whether the system is planar. + + Returns: + Always ``True``. + """ return True @property def segment_length(self) -> Array: - """Per-segment backbone lengths.""" + """Return the per-segment backbone lengths. + + Returns: + Array with shape ``(num_segments,)``. + """ return jnp.asarray(self.L) def _strain_basis_scaling_vector(self) -> Array: @@ -221,7 +388,21 @@ def _scaled_strain_basis(self, B_xi: Array) -> Array: return B_xi def cross_section_geometry(self, q: Array, s: Array) -> tuple[Array, Array]: - """Return the assumed solid circular cross-section and segment radius.""" + """Evaluate the circular cross-section at a backbone coordinate. + + Args: + q: Generalized coordinates. PCS cross-section geometry is + configuration-independent, but ``q`` is accepted for the common + continuum-robot interface. + s: Scalar global backbone coordinate. + + Returns: + A tuple containing the circular :class:`CrossSectionGeometry` tag + and a one-element array holding the selected segment radius. + + Raises: + ValueError: If ``s`` cannot be assigned to a segment. + """ segment_idx, _ = self.classify_segment(s) radius = jnp.asarray(self.r)[segment_idx] tag = jnp.asarray(CrossSectionGeometry.CIRCULAR, dtype=jnp.int32) @@ -241,7 +422,7 @@ def _set_params(self, params: PlanarPCSParams) -> None: ) # Add a zero for the orientation angle # Lengths of the segments - L = params.length + L = params.link.length L = jnp.asarray(L, dtype=jnp.float64) if L.shape != (self.num_segments,): raise ValueError( @@ -253,16 +434,11 @@ def _set_params(self, params: PlanarPCSParams) -> None: self.L_cum = L_cum # Radius of the segments - r = params.radius - r = jnp.asarray(r, dtype=jnp.float64) - if r.shape != (self.num_segments,): - raise ValueError( - f"radius must have shape ({self.num_segments},), got {r.shape}" - ) + r = jnp.asarray(params.link.cross_section.coefficients[:, 0], dtype=jnp.float64) self.r = r # Densities of the segments - rho = params.density + rho = params.link.density rho = jnp.asarray(rho, dtype=jnp.float64) if rho.shape != (self.num_segments,): raise ValueError( @@ -270,72 +446,9 @@ def _set_params(self, params: PlanarPCSParams) -> None: ) self.rho = rho - # Elastic modulus of the segments - E = params.young_modulus - E = jnp.asarray(E, dtype=jnp.float64) - if E.shape != (self.num_segments,): - raise ValueError( - f"young_modulus must have shape ({self.num_segments},), got {E.shape}" - ) - self.E = E - - # Shear modulus of the segments - G = params.shear_modulus - G = jnp.asarray(G, dtype=jnp.float64) - if G.shape != (self.num_segments,): - raise ValueError( - f"shear_modulus must have shape ({self.num_segments},), got {G.shape}" - ) - self.G = G - - def _explicit_damping_full_matrix(self, params: PlanarPCSParams) -> Array: - """Return the custom full damping matrix supplied in params.""" - if params.damping_matrix is None: - raise ValueError("damping_matrix is not set.") - expected_D_shape = (self.num_strains, self.num_strains) - D = jnp.asarray(params.damping_matrix, dtype=jnp.float64) - if D.shape != expected_D_shape: - raise ValueError( - f"damping_matrix must have shape {expected_D_shape}, got {D.shape}" - ) - return D - - def _material_damping_coefficients(self) -> Array: - """Return per-segment material damping coefficients in Pa*s (N*s/m^2).""" - params = self._current_body_params() - if params.material_damping_coefficient is None: - raise ValueError("material_damping_coefficient is not set.") - coefficient = jnp.asarray( - params.material_damping_coefficient, dtype=jnp.float64 - ) - if coefficient.ndim == 0: - return jnp.full((self.num_segments,), coefficient, dtype=jnp.float64) - if coefficient.shape != (self.num_segments,): - raise ValueError( - "material_damping_coefficient must be a scalar or have shape " - f"({self.num_segments},), got {coefficient.shape}." - ) - return coefficient - - def _compute_material_damping_full_matrix(self) -> Array: - """Compute damping using the solid-circle area and second moment.""" - coefficients = self._material_damping_coefficients() - - def damping_block(i: Array) -> Array: - I_i = self._local_second_moment_of_area(i) - A_i = self._local_cross_sectional_area(i) - damping_diag = jnp.stack([3.0 * I_i, 3.0 * A_i, A_i], axis=0) - return self.L[i] * coefficients[i] * jnp.diag(damping_diag) - - damping_blocks = vmap(damping_block)(jnp.arange(self.num_segments)) - return blk_diag(damping_blocks) - def _compute_damping_full_matrix(self) -> Array: - """Compute the current full damping matrix.""" - params = self._current_body_params() - if params.material_damping_coefficient is not None: - return self._compute_material_damping_full_matrix() - return self._explicit_damping_full_matrix(params) + """Assemble canonical per-link generalized damping blocks.""" + return blk_diag(self._current_body_params().link.damping) def _current_body_params(self) -> PlanarPCSParams: """Return the planar PCS body params, including typed actuated wrappers.""" @@ -354,23 +467,26 @@ def _with_planar_pcs_params( if not isinstance(params, PlanarPCSParams): raise TypeError("params must be a PlanarPCSParams instance.") params.validate() - if params.length.shape != current_params.length.shape: + if params.link.length.shape != current_params.link.length.shape: raise ValueError( "length shape changes the model structure; construct a new PlanarPCS." ) - if params.reference_strain.shape != current_params.reference_strain.shape: + if ( + params.link.reference_strain.shape + != current_params.link.reference_strain.shape + ): raise ValueError( "reference_strain shape changes the model structure; construct a new PlanarPCS." ) base_pose = jnp.asarray(params.base_pose, dtype=jnp.float64) gravity = jnp.asarray(params.gravity, dtype=jnp.float64) - segment_lengths = jnp.asarray(params.length, dtype=jnp.float64) - radius = jnp.asarray(params.radius, dtype=jnp.float64) - density = jnp.asarray(params.density, dtype=jnp.float64) - young_modulus = jnp.asarray(params.young_modulus, dtype=jnp.float64) - shear_modulus = jnp.asarray(params.shear_modulus, dtype=jnp.float64) - reference_strain = jnp.asarray(params.reference_strain, dtype=jnp.float64) + segment_lengths = jnp.asarray(params.link.length, dtype=jnp.float64) + radius = jnp.asarray( + params.link.cross_section.coefficients[:, 0], dtype=jnp.float64 + ) + density = jnp.asarray(params.link.density, dtype=jnp.float64) + reference_strain = jnp.asarray(params.link.reference_strain, dtype=jnp.float64) updated_self = eqx.tree_at( lambda m: ( @@ -381,8 +497,6 @@ def _with_planar_pcs_params( m.L_cum, m.r, m.rho, - m.E, - m.G, m.xi_ref, ), self, @@ -398,36 +512,133 @@ def _with_planar_pcs_params( ), radius, density, - young_modulus, - shear_modulus, reference_strain.reshape(self.num_strains), ), ) return updated_self._with_refreshed_precomputed_matrices() def with_params(self, params: PlanarPCSParams) -> "PlanarPCS": - """Return an updated copy with a full typed parameter object.""" + """Return a model copy using a complete parameter PyTree. + + Args: + params: Complete replacement parameters with the same number of + segments and reference-strain layout as this model. + + Returns: + A new PlanarPCS model with refreshed geometry, mass, + material-response, stiffness, and damping caches. The original + model is unchanged. + + Raises: + TypeError: If ``params`` is not a :class:`PlanarPCSParams`. + ValueError: If the replacement is invalid or changes the static + segment or strain layout. + """ return self._with_planar_pcs_params(params) - def update_params(self, **updates: Array) -> "PlanarPCS": - """Return an updated copy with selected typed parameter fields replaced.""" + def update_params(self, **updates: Any) -> "PlanarPCS": + """Return a copy with selected top-level parameter fields replaced. + + Args: + **updates: Fields of :class:`PlanarPCSParams` to replace, typically + ``link``, ``gravity``, or ``base_pose``. + + Returns: + A new validated PlanarPCS model containing the replacements. + + Raises: + TypeError: If an unknown field is supplied or a replacement has an + invalid type. + ValueError: If the result is invalid or changes the static segment + or strain layout. + """ if ( - "length" in updates - and jnp.asarray(updates["length"]).shape != self.params.length.shape + "link" in updates + and jnp.asarray(updates["link"].length).shape + != self.params.link.length.shape ): raise ValueError( "length shape changes the model structure; construct a new PlanarPCS." ) if ( - "reference_strain" in updates - and jnp.asarray(updates["reference_strain"]).shape - != self.params.reference_strain.shape + "link" in updates + and jnp.asarray(updates["link"].reference_strain).shape + != self.params.link.reference_strain.shape ): raise ValueError( "reference_strain shape changes the model structure; construct a new PlanarPCS." ) return self.with_params(self.params.replace(**updates)) + def update_link_params(self, **updates: Any) -> "PlanarPCS": + """Return a copy with selected continuum-link fields replaced. + + Args: + **updates: Fields of :class:`ContinuumLinkParams` to replace, such + as ``length``, ``density``, ``reference_strain``, + ``cross_section``, ``stiffness``, or ``damping``. + + Returns: + A new validated PlanarPCS model with refreshed dependent caches. + + Raises: + TypeError: If an unknown link field is supplied. + ValueError: If a replacement is invalid or changes the static + segment or strain layout. + """ + return self.with_params( + self.params.replace(link=self.params.link.replace(**updates)) + ) + + def link_matrices_from_material( + self, material: IsotropicMaterialParams + ) -> tuple[Array, Array]: + """Map isotropic material variables to canonical planar link matrices. + + Args: + material: Scalar or per-segment Young's modulus, shear modulus, and + material damping coefficient. Scalars are broadcast over all + segments. + + Returns: + A tuple ``(stiffness, damping)`` whose arrays both have shape + ``(num_segments, 3, 3)``. + + Raises: + ValueError: If a material field is not scalar or does not have + shape ``(num_segments,)``. + """ + material = material.broadcast(self.num_segments) + stiffness = ( + material.young_modulus[:, None, None] * self.young_stiffness_operator + + material.shear_modulus[:, None, None] * self.shear_stiffness_operator + ) + damping = ( + material.material_damping_coefficient[:, None, None] + * self.material_damping_operator + ) + return stiffness, damping + + def with_isotropic_material(self, material: IsotropicMaterialParams) -> Self: + """Return a copy whose link matrices are built from isotropic material. + + The supplied material PyTree remains caller-owned and is not stored on + the model. + + Args: + material: Scalar or per-segment isotropic material variables. + + Returns: + A new PlanarPCS model containing the generated canonical link + stiffness and damping matrices. + + Raises: + ValueError: If a material field cannot be broadcast to one value + per segment. + """ + stiffness, damping = self.link_matrices_from_material(material) + return self.update_link_params(stiffness=stiffness, damping=damping) + def _precomputed_matrices(self) -> tuple[Array, Array, Array, Array, Array]: """Compute state-independent matrices cached by the model.""" M_segments = vmap(self._compute_local_mass_matrix)( @@ -440,8 +651,17 @@ def _precomputed_matrices(self) -> tuple[Array, Array, Array, Array, Array]: return M_segments, K_full, K_active, D_full, D_active def precompute(self) -> None: - """Refresh state-independent matrices cached by the model.""" + """Refresh state-independent matrices cached by the model. + + Returns: + ``None``. Material unit-response operators and the segment, full, + and active-coordinate matrices are replaced in place. + """ object.__setattr__(self, "B_xi", self._scaled_strain_basis(self.B_xi_unscaled)) + young_operator, shear_operator, damping_operator = self._material_operators() + object.__setattr__(self, "young_stiffness_operator", young_operator) + object.__setattr__(self, "shear_stiffness_operator", shear_operator) + object.__setattr__(self, "material_damping_operator", damping_operator) ( M_segments, K_full, @@ -459,6 +679,18 @@ def _with_refreshed_precomputed_matrices(self) -> "PlanarPCS": """Return a copy with cached state-independent matrices refreshed.""" B_xi = self._scaled_strain_basis(self.B_xi_unscaled) updated_self = eqx.tree_at(lambda m: m.B_xi, self, B_xi) + young_operator, shear_operator, damping_operator = ( + updated_self._material_operators() + ) + updated_self = eqx.tree_at( + lambda m: ( + m.young_stiffness_operator, + m.shear_stiffness_operator, + m.material_damping_operator, + ), + updated_self, + (young_operator, shear_operator, damping_operator), + ) ( M_segments, K_full, @@ -1469,8 +1701,17 @@ def skip_branch( def jacobian_and_arc_length_derivative_bodyframe( self, q: Array, s: Array ) -> tuple[Array, Array]: - """ - Compute the body-frame Jacobian and its arc-length derivative at ``s``. + """Compute a body-frame Jacobian and its arc-length derivative. + + Args: + q: Active generalized strains with shape + ``(num_active_strains,)``. + s: Scalar global backbone coordinate. + + Returns: + A tuple ``(J, J_s)`` containing the body-frame Jacobian and its + derivative with respect to global arc length. Both arrays have + shape ``(3, num_active_strains)``. """ xi = self.strain(q).reshape(self.num_segments, 3) segment_idx, s_local = self.classify_segment(s) @@ -1553,8 +1794,16 @@ def skip_branch( @eqx.filter_jit def jacobian_arc_length_derivative_bodyframe(self, q: Array, s: Array) -> Array: - """ - Compute the arc-length derivative of the body-frame Jacobian at ``s``. + """Compute the body-frame Jacobian derivative with respect to arc length. + + Args: + q: Active generalized strains with shape + ``(num_active_strains,)``. + s: Scalar global backbone coordinate. + + Returns: + Arc-length derivative with shape + ``(3, num_active_strains)``. """ _, Js = self.jacobian_and_arc_length_derivative_bodyframe(q, s) return Js @@ -1606,8 +1855,17 @@ def jacobian_inertialframe(self, q: Array, s: Array) -> Array: def jacobian_and_arc_length_derivative_inertialframe( self, q: Array, s: Array ) -> tuple[Array, Array]: - """ - Compute the inertial-frame Jacobian and its arc-length derivative at ``s``. + """Compute an inertial-frame Jacobian and its arc-length derivative. + + Args: + q: Active generalized strains with shape + ``(num_active_strains,)``. + s: Scalar global backbone coordinate. + + Returns: + A tuple ``(J, J_s)`` containing the inertial-frame Jacobian and its + derivative with respect to global arc length. Both arrays have + shape ``(3, num_active_strains)``. """ chi, J_local, Js_local = ( self._jacobian_and_arc_length_derivative_bodyframe_with_pose(q, s) @@ -1626,8 +1884,16 @@ def jacobian_and_arc_length_derivative_inertialframe( @eqx.filter_jit def jacobian_arc_length_derivative_inertialframe(self, q: Array, s: Array) -> Array: - """ - Compute the arc-length derivative of the inertial-frame Jacobian at ``s``. + """Compute the inertial Jacobian derivative with respect to arc length. + + Args: + q: Active generalized strains with shape + ``(num_active_strains,)``. + s: Scalar global backbone coordinate. + + Returns: + Arc-length derivative with shape + ``(3, num_active_strains)``. """ chi, J_local, Js_local = ( self._jacobian_and_arc_length_derivative_bodyframe_with_pose(q, s) @@ -2047,7 +2313,17 @@ def _jacobian_and_arc_length_derivative( @eqx.filter_jit def jacobian_batched(self, q: Array, s_ps: Array) -> Array: - """Compute inertial-frame Jacobians at multiple arc-length positions.""" + """Compute inertial-frame Jacobians at multiple arc-length positions. + + Args: + q: Active generalized strains with shape + ``(num_active_strains,)``. + s_ps: Backbone coordinates with shape ``(num_points,)``. + + Returns: + Inertial-frame Jacobians with shape + ``(num_points, 3, num_active_strains)``. + """ return self.jacobian_inertialframe_batched(q, s_ps) @eqx.filter_jit @@ -2061,7 +2337,19 @@ def _jacobian_and_time_derivative( def jacobian_and_time_derivative_batched( self, q: Array, qd: Array, s_ps: Array ) -> tuple[Array, Array]: - """Compute inertial-frame Jacobians and time derivatives at multiple arc-length positions.""" + """Compute batched inertial Jacobians and their time derivatives. + + Args: + q: Active generalized strains with shape + ``(num_active_strains,)``. + qd: Active generalized strain rates with shape + ``(num_active_strains,)``. + s_ps: Backbone coordinates with shape ``(num_points,)``. + + Returns: + A tuple ``(J, J_dot)`` whose arrays both have shape + ``(num_points, 3, num_active_strains)``. + """ return self.jacobian_and_time_derivative_inertialframe_batched(q, qd, s_ps) # ========================================== @@ -2340,43 +2628,25 @@ def _gravitational_force(self, q: Array) -> Array: return G - @eqx.filter_jit - def _local_stiffness_matrix(self, i: Array) -> Array: - """ - Compute local stiffness for a planar rod longitudinally aligned with local x. - - The diagonal constitutive terms use the solid circular cross-section's - area and transverse second moment. - - Args: - i (Array): index of the segment - - Returns: - S_i (Array): Local stiffness matrix of shape (3, 3) for the i-th segment. - """ - I_i = self._local_second_moment_of_area(i) # Second moment of area - A_i = self._local_cross_sectional_area(i) # Cross-sectional area - - S_i = self.L[i] * jnp.diag( - jnp.stack( - [ - I_i * self.E[i], # bending Z - A_i * self.E[i], # axial X - A_i * self.G[i], # shear Y - ], - axis=0, - ) + def _material_operators(self) -> tuple[Array, Array, Array]: + """Return unit Young, shear, and material-damping link matrices.""" + area = jnp.pi * self.r**2 + moment = jnp.pi * self.r**4 / 4.0 + zeros = jnp.zeros_like(self.L) + young = self.L[:, None, None] * vmap(jnp.diag)( + jnp.stack([moment, area, zeros], axis=1) ) - - return S_i + shear = self.L[:, None, None] * vmap(jnp.diag)( + jnp.stack([zeros, zeros, area], axis=1) + ) + damping = self.L[:, None, None] * vmap(jnp.diag)( + jnp.stack([3.0 * moment, 3.0 * area, area], axis=1) + ) + return young, shear, damping def _compute_stiffness_full_matrix(self) -> Array: - """Compute the uncached full stiffness matrix from current parameters.""" - # stiffness matrix of shape (num_segments, 3, 3) - S_sms = vmap(self._local_stiffness_matrix)(jnp.arange(self.num_segments)) - - # we define the elastic matrix of shape (num_strains, num_strains) as K(xi) = K @ xi where K is equal to - return blk_diag(S_sms) + """Assemble canonical per-link generalized stiffness blocks.""" + return blk_diag(self._current_body_params().link.stiffness) @eqx.filter_jit def _stiffness(self, formulate_in_strain_space: bool = False) -> Array: diff --git a/src/soromox/systems/pcs/structures.py b/src/soromox/systems/pcs/structures.py index c650d71e..e1020731 100644 --- a/src/soromox/systems/pcs/structures.py +++ b/src/soromox/systems/pcs/structures.py @@ -44,7 +44,15 @@ def _rigid_segment_selector(value: Any) -> tuple[bool, ...] | None: class PCSStructure(eqx.Module): - """Static PCS layout that determines JAX compilation structure.""" + """Static spatial PCS layout that determines JAX compilation structure. + + Attributes: + num_gauss_points: Number of Gauss points per constant-strain link. + strain_selector: Optional boolean selector with six entries per link. + ``None`` activates every strain coordinate. + scale_rotational_basis_by_length: Whether rotational strain coordinates + are normalized by their link length. + """ num_gauss_points: int = eqx.field(static=True, default=5) strain_selector: Array | None = None @@ -52,7 +60,15 @@ class PCSStructure(eqx.Module): class PlanarPCSStructure(eqx.Module): - """Static planar PCS layout.""" + """Static planar PCS layout. + + Attributes: + num_gauss_points: Number of Gauss points per constant-strain link. + strain_selector: Optional boolean selector with three entries per link. + ``None`` activates every strain coordinate. + scale_rotational_basis_by_length: Whether the rotational strain + coordinate is normalized by its link length. + """ num_gauss_points: int = eqx.field(static=True, default=5) strain_selector: Array | None = None @@ -72,6 +88,12 @@ class ISupportStructure(PCSStructure): segment. When omitted, segment types alternate from a rigid segment at index zero. If ``strain_selector`` is provided, it is interpreted on the expanded PCS layout; rigid-segment strains are always deactivated. + + Attributes: + pcs_segment_counts: Optional number of PCS subdivisions for each + pneumatic segment. A one-element tuple is broadcast. + rigid_segment_selector: Optional physical-segment mask in which + ``True`` denotes a rigid connector. """ pcs_segment_counts: tuple[int, ...] | None = eqx.field( diff --git a/src/soromox/systems/pendulum/pendulum.py b/src/soromox/systems/pendulum/pendulum.py index e684d79e..5d74ee41 100644 --- a/src/soromox/systems/pendulum/pendulum.py +++ b/src/soromox/systems/pendulum/pendulum.py @@ -8,8 +8,9 @@ from jax import numpy as jnp from soromox.actuation.core import Actuator, PassiveElement +from soromox.systems.components import CrossSectionGeometry from soromox.systems.pendulum.params import PendulumParams -from soromox.systems.soft_robot import CrossSectionGeometry, SoftRobot +from soromox.systems.soft_robot import SoftRobot class Pendulum(SoftRobot): diff --git a/src/soromox/systems/soft_robot.py b/src/soromox/systems/soft_robot.py index 755b1cac..c50166a2 100644 --- a/src/soromox/systems/soft_robot.py +++ b/src/soromox/systems/soft_robot.py @@ -1,10 +1,8 @@ __all__ = [ - "CrossSectionGeometry", "SoftRobot", ] from abc import abstractmethod -from enum import IntEnum from typing import Any import equinox as eqx @@ -26,12 +24,6 @@ from soromox.utils.geometry import poses -class CrossSectionGeometry(IntEnum): - CIRCULAR = 0 - RECTANGULAR = 1 - ELLIPTICAL = 2 - - class SoftRobot(DynamicalSystem): """ Abstract base class for soft robot systems. diff --git a/tests/actuation/test_threadlike.py b/tests/actuation/test_threadlike.py index 2505a1f2..fb570518 100644 --- a/tests/actuation/test_threadlike.py +++ b/tests/actuation/test_threadlike.py @@ -22,8 +22,16 @@ ThreadlikeImpedance, ThreadlikeRouting, ) -from soromox.systems import GVS, PCS, CrossSectionGeometry, PCSStructure, PlanarPCS -from soromox.systems.gvs import GVSSegment, JointSpec, LinkSpec, StrainBasisSpec +from soromox.systems import ( + GVS, + PCS, + GVSSegment, + JointSpec, + LinkSpec, + PCSStructure, + PlanarPCS, + StrainBasisSpec, +) from soromox.systems.pcs import PlanarPCSStructure @@ -67,21 +75,20 @@ def _planar_pcs(*, actuators=None, passive_elements=()): def _gvs(*, actuators=None, passive_elements=()): segment = GVSSegment( - link=LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=3e5, - nu=0.45, - rho=1300.0, - eta=1e4, - L=0.1, - r_i=0.015, - r_f=0.015, + link=LinkSpec.circular( + young_modulus=3e5, + shear_modulus=3e5 / 2.9, + density=1300.0, + material_damping_coefficient=1e4, + length=0.1, + radius=0.015, + reference_strain=[0, 0, 0, 1, 0, 0], ), joint=JointSpec(type="fixed"), basis=StrainBasisSpec( type="monomial", - active=[1, 1, 1, 1, 0, 0], - orders=[0, 0, 0, 0, 0, 0], + strain_selector=[1, 1, 1, 1, 0, 0], + basis_order=[0, 0, 0, 0, 0, 0], ), num_gauss_points=5, ) @@ -408,8 +415,10 @@ def test_actuator_nested_update_and_topology_rejection(): with pytest.raises(ValueError, match="topology"): params.transmission.routing.replace(end_segment_index=(1,)) - body_updated = robot.update_params(young_modulus=jnp.array([2e3])) - assert_allclose(body_updated.params.young_modulus, jnp.array([2e3])) + body_updated = robot.update_link_params(stiffness=2.0 * robot.params.link.stiffness) + assert_allclose( + body_updated.params.link.stiffness, 2.0 * robot.params.link.stiffness + ) assert_allclose( body_updated.actuators[0].params.transmission.routing.intercept, robot.actuators[0].params.transmission.routing.intercept, diff --git a/tests/rendering/test_base_renderer.py b/tests/rendering/test_base_renderer.py index b814bc1e..238c4a7e 100644 --- a/tests/rendering/test_base_renderer.py +++ b/tests/rendering/test_base_renderer.py @@ -17,7 +17,7 @@ from soromox.rendering.color_config import ActuatorStyleConfig from soromox.rendering.matplotlib_renderer import MatplotlibRenderer from soromox.rendering.opencv_planar_renderer import OpenCVPlanarRenderer -from soromox.systems.soft_robot import CrossSectionGeometry +from soromox.systems.components import CrossSectionGeometry from soromox.utils.geometry import poses diff --git a/tests/rendering/test_isupport_viser_renderer.py b/tests/rendering/test_isupport_viser_renderer.py index e002d9a1..7fb466bf 100644 --- a/tests/rendering/test_isupport_viser_renderer.py +++ b/tests/rendering/test_isupport_viser_renderer.py @@ -16,7 +16,13 @@ from soromox.rendering.color_config import validate_rgb from soromox.rendering.isupport.viser_renderer import ISupportLiveModeController from soromox.rendering.viser_renderer import SceneHandles -from soromox.systems import ISupport, ISupportParams, ISupportStructure +from soromox.systems import ( + PCS, + ISupport, + ISupportParams, + ISupportStructure, + LinkSpec, +) class _FakeHandle: @@ -98,19 +104,22 @@ def _make_robot(*, connectors: bool = True) -> ISupport: radii = jnp.array([0.03, 0.03]) densities = jnp.array([1000.0, 1000.0]) num_physical_segments = len(rigid_segment_selector) + links = [ + LinkSpec.circular( + length=float(lengths[index]), + radius=float(radii[index]), + density=float(densities[index]), + young_modulus=2.0e3, + shear_modulus=1.0e3, + damping=1.0e-3 * jnp.eye(6), + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + ) + for index in range(num_physical_segments) + ] params = ISupportParams( + link=PCS.params_from_links(links).link, base_pose=jnp.array([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), - length=lengths, - radius=radii, - density=densities, gravity=jnp.array([0.0, 0.0, -9.81]), - young_modulus=2.0e3 * jnp.ones((num_physical_segments,)), - shear_modulus=1.0e3 * jnp.ones((num_physical_segments,)), - damping_matrix=1.0e-3 * jnp.eye(6 * num_physical_segments), - reference_strain=jnp.tile( - jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), - num_physical_segments, - ), chamber_inner_radius=jnp.array([0.005, 0.005]), chamber_outer_radius=jnp.array([0.00779, 0.00779]), chamber_distance=jnp.array([0.020, 0.020]), diff --git a/tests/rendering/test_open3d_material_frames.py b/tests/rendering/test_open3d_material_frames.py index ca506c25..3c8fd87e 100644 --- a/tests/rendering/test_open3d_material_frames.py +++ b/tests/rendering/test_open3d_material_frames.py @@ -9,7 +9,7 @@ _make_swept_cross_section_segment, _swept_segment_vertices, ) -from soromox.systems.soft_robot import CrossSectionGeometry # noqa: E402 +from soromox.systems.components import CrossSectionGeometry # noqa: E402 def test_swept_cylinder_rings_follow_material_frame_not_curve_chord(): diff --git a/tests/system_param_builders.py b/tests/system_param_builders.py index cfc315c5..14969327 100644 --- a/tests/system_param_builders.py +++ b/tests/system_param_builders.py @@ -1,9 +1,12 @@ +import jax import jax.numpy as jnp from jax import Array from soromox.systems import ( GVS, ArticulatedSoftRobotParams, + ContinuumLinkParams, + CrossSectionParams, GVSParams, GVSStructure, PCSParams, @@ -45,16 +48,45 @@ def pcs_params( reference_strain = jnp.tile( jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), num_segments ) + radius = jnp.asarray(radius) + young_modulus = jnp.asarray(young_modulus) + shear_modulus = jnp.asarray(shear_modulus) + area = jnp.pi * radius**2 + transverse = jnp.pi * radius**4 / 4.0 + polar = 2.0 * transverse + stiffness = length[:, None, None] * jax.vmap(jnp.diag)( + jnp.stack( + [ + shear_modulus * polar, + young_modulus * transverse, + young_modulus * transverse, + young_modulus * area, + shear_modulus * area, + shear_modulus * area, + ], + axis=1, + ) + ) + damping_matrix = jnp.asarray(damping_matrix) + damping = jnp.stack( + [ + damping_matrix[6 * i : 6 * (i + 1), 6 * i : 6 * (i + 1)] + for i in range(num_segments) + ] + ) + if not bool(jnp.allclose(damping_matrix, jax.scipy.linalg.block_diag(*damping))): + raise ValueError("PCS cross-link damping coupling is no longer supported.") return PCSParams( - length=length, - radius=jnp.asarray(radius), - density=jnp.asarray(density), - young_modulus=jnp.asarray(young_modulus), - shear_modulus=jnp.asarray(shear_modulus), - damping_matrix=jnp.asarray(damping_matrix), gravity=jnp.asarray(gravity), base_pose=jnp.asarray(base_pose), - reference_strain=jnp.asarray(reference_strain), + link=ContinuumLinkParams( + length=length, + density=jnp.asarray(density), + reference_strain=jnp.asarray(reference_strain).reshape(num_segments, 6), + cross_section=CrossSectionParams(coefficients=radius[:, None]), + stiffness=stiffness, + damping=damping, + ), ) @@ -76,16 +108,37 @@ def planar_pcs_params( reference_strain = jnp.tile(jnp.array([0.0, 1.0, 0.0]), num_segments) if base_pose is None: base_pose = planar_base_pose(jnp.pi / 2) + radius = jnp.asarray(radius) + young_modulus = jnp.asarray(young_modulus) + shear_modulus = jnp.asarray(shear_modulus) + area = jnp.pi * radius**2 + moment = jnp.pi * radius**4 / 4.0 + stiffness = length[:, None, None] * jax.vmap(jnp.diag)( + jnp.stack( + [young_modulus * moment, young_modulus * area, shear_modulus * area], + axis=1, + ) + ) + damping_matrix = jnp.asarray(damping_matrix) + damping = jnp.stack( + [ + damping_matrix[3 * i : 3 * (i + 1), 3 * i : 3 * (i + 1)] + for i in range(num_segments) + ] + ) + if not bool(jnp.allclose(damping_matrix, jax.scipy.linalg.block_diag(*damping))): + raise ValueError("PCS cross-link damping coupling is no longer supported.") return PlanarPCSParams( - length=length, - radius=jnp.asarray(radius), - density=jnp.asarray(density), - young_modulus=jnp.asarray(young_modulus), - shear_modulus=jnp.asarray(shear_modulus), - damping_matrix=jnp.asarray(damping_matrix), gravity=jnp.asarray(gravity), base_pose=jnp.asarray(base_pose), - reference_strain=jnp.asarray(reference_strain), + link=ContinuumLinkParams( + length=length, + density=jnp.asarray(density), + reference_strain=jnp.asarray(reference_strain).reshape(num_segments, 3), + cross_section=CrossSectionParams(coefficients=radius[:, None]), + stiffness=stiffness, + damping=damping, + ), ) diff --git a/tests/systems/test_gvs.py b/tests/systems/test_gvs.py index 52795d13..d6ee1b9a 100644 --- a/tests/systems/test_gvs.py +++ b/tests/systems/test_gvs.py @@ -9,8 +9,18 @@ spatial_base_pose, ) -from soromox.systems import CrossSectionGeometry, GVS, PCS, PCSParams, PCSStructure -from soromox.systems.gvs import GVSSegment, JointSpec, LinkSpec, StrainBasisSpec +from soromox.systems import ( + GVS, + PCS, + CrossSectionGeometry, + GVSSegment, + IsotropicMaterialParams, + JointSpec, + LinearProfile, + LinkSpec, + PCSStructure, + StrainBasisSpec, +) from soromox.utils.lie_algebra import se3 from soromox.utils.tolerance import Tolerance @@ -50,22 +60,20 @@ def build_matched_gvs_pcs( # GVS definition: constant strain along each link, all 6 strain components enabled segments = [ GVSSegment( - link=LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=float(E[i]), - nu=float(nu[i]), - rho=float(rhos[i]), - eta=material_damping_coefficient, - L=float(Ls[i]), - r_i=float(rs[i]), - r_f=float(rs[i]), + link=LinkSpec.circular( + young_modulus=float(E[i]), + shear_modulus=float(Gpcs[i]), + density=float(rhos[i]), + material_damping_coefficient=material_damping_coefficient, + length=float(Ls[i]), + radius=float(rs[i]), + reference_strain=[0, 0, 0, 1, 0, 0], ), joint=JointSpec(type="fixed"), basis=StrainBasisSpec( type="monomial", - active=[1, 1, 1, 1, 1, 1], - orders=[0, 0, 0, 0, 0, 0], - xi_ref=[0, 0, 0, 1, 0, 0], + strain_selector=[1, 1, 1, 1, 1, 1], + basis_order=[0, 0, 0, 0, 0, 0], ), num_gauss_points=n_gauss, ) @@ -81,26 +89,14 @@ def build_matched_gvs_pcs( ) robot_gvs = GVS(params=gvs_params, structure=gvs_structure) - # PCS definition with identical geometry and material params - params = PCSParams( + # PCS definition through the same shared link specifications. This exercises + # the common isotropic-material mapping rather than an explicit matrix path. + robot_pcs = PCS.from_links( + [segment.link for segment in segments], base_pose=spatial_base_pose(), - length=Ls, - radius=rs, - density=rhos, gravity=g, - young_modulus=E, - shear_modulus=Gpcs, - material_damping_coefficient=jnp.asarray( - material_damping_coefficient, dtype=jnp.float64 - ), - reference_strain=jnp.tile( - jnp.array([0, 0, 0, 1, 0, 0]), (num_segments, 1) - ).reshape(6 * num_segments), - ) - robot_pcs = PCS( - params=params, structure=PCSStructure( - num_gauss_points=5, + num_gauss_points=n_gauss, strain_selector=jnp.ones((6 * num_segments,), dtype=bool), ), ) @@ -110,22 +106,20 @@ def build_matched_gvs_pcs( def test_params_from_segments_stores_resolved_max_dof(): segment = GVSSegment( - link=LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=1e6, - nu=0.45, - rho=1000.0, - eta=0.0, - L=0.2, - r_i=0.02, - r_f=0.02, + link=LinkSpec.circular( + young_modulus=1e6, + shear_modulus=1e6 / 2.9, + density=1000.0, + material_damping_coefficient=0.0, + length=0.2, + radius=0.02, + reference_strain=[0, 0, 0, 1, 0, 0], ), joint=JointSpec(type="fixed"), basis=StrainBasisSpec( type="monomial", - active=[1, 1, 1, 1, 1, 1], - orders=[0, 0, 0, 0, 0, 0], - xi_ref=[0, 0, 0, 1, 0, 0], + strain_selector=[1, 1, 1, 1, 1, 1], + basis_order=[0, 0, 0, 0, 0, 0], ), num_gauss_points=5, ) @@ -134,33 +128,32 @@ def test_params_from_segments_stores_resolved_max_dof(): [segment], gravity=jnp.array([0.0, 0.0, -9.81]) ) - assert structure.max_dof == params.joint_stiffness.shape[1] == 6 + assert structure.max_dof == params.joint.stiffness.shape[1] == 6 - oversized = params.replace(joint_stiffness=jnp.zeros((1, 7, 7))) - with pytest.raises(ValueError, match="joint_stiffness"): - oversized.validate_against_structure(structure) - with pytest.raises(ValueError, match="joint_stiffness"): - GVS(params=oversized, structure=structure) + with pytest.raises(ValueError, match="joint"): + params.replace( + joint=params.joint.replace( + stiffness=jnp.zeros((1, 7, 7)), damping=jnp.zeros((1, 7, 7)) + ) + ) def test_params_from_segments_uses_spatial_environment_defaults(): segment = GVSSegment( - link=LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=1e6, - nu=0.45, - rho=1000.0, - eta=0.0, - L=0.2, - r_i=0.02, - r_f=0.02, + link=LinkSpec.circular( + young_modulus=1e6, + shear_modulus=1e6 / 2.9, + density=1000.0, + material_damping_coefficient=0.0, + length=0.2, + radius=0.02, + reference_strain=[0, 0, 0, 1, 0, 0], ), joint=JointSpec(type="fixed"), basis=StrainBasisSpec( type="monomial", - active=[1, 1, 1, 1, 1, 1], - orders=[0, 0, 0, 0, 0, 0], - xi_ref=[0, 0, 0, 1, 0, 0], + strain_selector=[1, 1, 1, 1, 1, 1], + basis_order=[0, 0, 0, 0, 0, 0], ), num_gauss_points=5, ) @@ -191,15 +184,14 @@ def _circular_link(idx: int) -> LinkSpec: scale = 1.0 + 0.04 * repeat r_i = 0.015 + 0.0008 * repeat r_f = r_i + 0.003 + 0.0004 * repeat - return LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=1.2e6, - nu=0.45, - rho=950.0, - eta=5.0, - L=float(0.25 * scale), - r_i=float(r_i), - r_f=float(r_f), + return LinkSpec.circular( + young_modulus=1.2e6, + shear_modulus=1.2e6 / (2 * 1.45), + density=950.0, + material_damping_coefficient=5.0, + length=float(0.25 * scale), + radius=LinearProfile(float(r_i), float(r_f)), + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], ) def _rectangular_link(idx: int) -> LinkSpec: @@ -209,17 +201,15 @@ def _rectangular_link(idx: int) -> LinkSpec: h_f = max(0.022, h_i * 0.9) w_i = 0.02 + 0.0008 * repeat w_f = max(0.016, w_i * 0.88) - return LinkSpec( - cross_section_geometry=CrossSectionGeometry.RECTANGULAR, - E=9.5e5, - nu=0.38, - rho=1025.0, - eta=4.0, - L=float(0.18 * scale), - h_i=float(h_i), - h_f=float(h_f), - w_i=float(w_i), - w_f=float(w_f), + return LinkSpec.rectangular( + young_modulus=9.5e5, + shear_modulus=9.5e5 / (2 * 1.38), + density=1025.0, + material_damping_coefficient=4.0, + length=float(0.18 * scale), + height=LinearProfile(float(h_i), float(h_f)), + width=LinearProfile(float(w_i), float(w_f)), + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], ) def _elliptical_link(idx: int) -> LinkSpec: @@ -229,17 +219,16 @@ def _elliptical_link(idx: int) -> LinkSpec: a_f = max(0.014, a_i * 0.92) b_i = 0.015 + 0.0007 * repeat b_f = b_i * 1.05 - return LinkSpec( - cross_section_geometry=CrossSectionGeometry.ELLIPTICAL, - E=8.0e5, - nu=0.4, - rho=980.0, - eta=3.5, - L=float(0.22 * scale), - a_i=float(a_i), - a_f=float(a_f), - b_i=float(b_i), - b_f=float(b_f), + xi_sigma = max(0.7, 0.9 - 0.05 * repeat) + return LinkSpec.elliptical( + young_modulus=8.0e5, + shear_modulus=8.0e5 / (2 * 1.4), + density=980.0, + material_damping_coefficient=3.5, + length=float(0.22 * scale), + semi_major=LinearProfile(float(a_i), float(a_f)), + semi_minor=LinearProfile(float(b_i), float(b_f)), + reference_strain=[0.2 + 0.02 * repeat, 0.0, 0.0, xi_sigma, 0.0, 0.0], ) def _revolute_joint(idx: int) -> JointSpec: @@ -262,27 +251,31 @@ def _monomial_basis(idx: int) -> StrainBasisSpec: extra = repeat % 2 return StrainBasisSpec( type="monomial", - active=[1, 1, 0, 1, 0, 0], - orders=[2 + extra, 1 + (idx % 2), 0, 2 + ((idx + repeat) % 2), 0, 0], - xi_ref=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + strain_selector=[1, 1, 0, 1, 0, 0], + basis_order=[2 + extra, 1 + (idx % 2), 0, 2 + ((idx + repeat) % 2), 0, 0], ) def _legendre_basis(idx: int) -> StrainBasisSpec: repeat = idx // pattern_count return StrainBasisSpec( type="legendre", - active=[0, 1, 1, 0, 1, 0], - orders=[0, 2 + (repeat % 2), 1 + ((idx + 1) % 2), 0, 1 + (repeat % 3), 0], - xi_ref=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + strain_selector=[0, 1, 1, 0, 1, 0], + basis_order=[ + 0, + 2 + (repeat % 2), + 1 + ((idx + 1) % 2), + 0, + 1 + (repeat % 3), + 0, + ], ) def _fourier_basis(idx: int) -> StrainBasisSpec: repeat = idx // pattern_count - xi_sigma = max(0.7, 0.9 - 0.05 * repeat) return StrainBasisSpec( type="fourier", - active=[1, 0, 1, 1, 0, 1], - orders=[ + strain_selector=[1, 0, 1, 1, 0, 1], + basis_order=[ 1 + (idx % 2), 0, 2 + (repeat % 2), @@ -290,7 +283,6 @@ def _fourier_basis(idx: int) -> StrainBasisSpec: 0, 1 + ((repeat + 1) % 2), ], - xi_ref=[0.2 + 0.02 * repeat, 0.0, 0.0, xi_sigma, 0.0, 0.0], ) def _monomial_gauss(idx: int) -> int: @@ -354,22 +346,20 @@ def build_constant_strain_gvs( segments = [ GVSSegment( - link=LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=1e6, - nu=0.5, - rho=1000.0, - eta=0.0, - L=segment_length, - r_i=0.02, - r_f=0.02, + link=LinkSpec.circular( + young_modulus=1e6, + shear_modulus=1e6 / 3.0, + density=1000.0, + material_damping_coefficient=0.0, + length=segment_length, + radius=0.02, + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], ), joint=JointSpec(type="fixed"), basis=StrainBasisSpec( type="monomial", - active=[int(active) for active in selector_per_segment], - orders=[0, 0, 0, 0, 0, 0], - xi_ref=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + strain_selector=[int(active) for active in selector_per_segment], + basis_order=[0, 0, 0, 0, 0, 0], ), num_gauss_points=5, ) @@ -390,36 +380,38 @@ def test_gvs_segment_factories_match_explicit_constructor() -> None: segments = [ GVSSegment( link=LinkSpec.circular( - E=1.0e6, - nu=0.45, - rho=1000.0, - eta=1.0, - L=0.2, - r=0.02, + young_modulus=1.0e6, + shear_modulus=1.0e6 / 2.9, + density=1000.0, + material_damping_coefficient=1.0, + length=0.2, + radius=0.02, + reference_strain=[0, 0, 0, 1, 0, 0], ), joint=JointSpec(type="fixed"), basis=StrainBasisSpec( type="monomial", - active=[1, 1, 1, 1, 0, 0], - orders=[0, 0, 0, 0, 0, 0], + strain_selector=[1, 1, 1, 1, 0, 0], + basis_order=[0, 0, 0, 0, 0, 0], ), num_gauss_points=5, ), GVSSegment( link=LinkSpec.rectangular( - E=9.0e5, - nu=0.4, - rho=950.0, - eta=2.0, - L=0.15, - h=0.03, - w=0.02, + young_modulus=9.0e5, + shear_modulus=9.0e5 / 2.8, + density=950.0, + material_damping_coefficient=2.0, + length=0.15, + height=0.03, + width=0.02, + reference_strain=[0, 0, 0, 1, 0, 0], ), joint=JointSpec(type="revolute", axis="z", stiffness=jnp.array([[0.3]])), basis=StrainBasisSpec( type="legendre", - active=[0, 1, 1, 0, 0, 0], - orders=[0, 1, 1, 0, 0, 0], + strain_selector=[0, 1, 1, 0, 0, 0], + basis_order=[0, 1, 1, 0, 0, 0], ), num_gauss_points=6, ), @@ -441,7 +433,7 @@ def test_gvs_segment_factories_match_explicit_constructor() -> None: ) assert_allclose(factory.params.link.length, params.link.length) - assert_allclose(factory.joint_stiffness, explicit.joint_stiffness) + assert_allclose(factory.params.joint.stiffness, explicit.params.joint.stiffness) assert_allclose( factory.cross_section_geometry_index, explicit.cross_section_geometry_index ) @@ -454,19 +446,19 @@ def test_gvs_segment_factories_match_explicit_constructor() -> None: def test_gvs_structure_contains_only_static_segment_choices() -> None: segment = GVSSegment( link=LinkSpec.circular( - E=1.0e6, - nu=0.45, - rho=1000.0, - eta=1.0, - L=0.2, - r=0.02, + young_modulus=1.0e6, + shear_modulus=1.0e6 / 2.9, + density=1000.0, + material_damping_coefficient=1.0, + length=0.2, + radius=0.02, + reference_strain=[0, 0, 0, 1, 0, 0], ), joint=JointSpec(type="revolute", axis="z", stiffness=jnp.array([[0.3]])), basis=StrainBasisSpec( type="monomial", - active=[1, 1, 0, 0, 0, 0], - orders=[0, 1, 0, 0, 0, 0], - xi_ref=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + strain_selector=[1, 1, 0, 0, 0, 0], + basis_order=[0, 1, 0, 0, 0, 0], ), num_gauss_points=5, ) @@ -483,11 +475,10 @@ def test_gvs_structure_contains_only_static_segment_choices() -> None: assert stored_segment.joint.type == "revolute" assert stored_segment.joint.axis == "z" assert not hasattr(stored_segment.joint, "stiffness") - assert stored_segment.basis.active == (1, 1, 0, 0, 0, 0) - assert stored_segment.basis.orders == (0, 1, 0, 0, 0, 0) + assert stored_segment.basis.strain_selector == (1, 1, 0, 0, 0, 0) + assert stored_segment.basis.basis_order == (0, 1, 0, 0, 0, 0) assert not hasattr(stored_segment.basis, "xi_ref") - assert_allclose(params.link.young_modulus, jnp.array([1.0e6])) - assert_allclose(params.joint_stiffness[0, 0, 0], 0.3) + assert_allclose(params.joint.stiffness[0, 0, 0], 0.3) def sample_arc_lengths(robot: GVS) -> jnp.ndarray: @@ -1609,14 +1600,18 @@ def test_cached_constant_matrices_refresh_after_update_params() -> None: max_dof=6, ) - updated = robot.update_params( - link=robot.params.link.replace( + section = robot.params.link.cross_section.replace( + coefficients=0.022 * jnp.ones_like(robot.params.link.cross_section.coefficients) + ) + geometry_updated = robot.update_link_params( + density=900.0 * jnp.ones_like(robot.segment_length), + cross_section=section, + ) + updated = geometry_updated.with_isotropic_material( + IsotropicMaterialParams( young_modulus=1.25e6 * jnp.ones_like(robot.segment_length), - poisson_ratio=0.45 * jnp.ones_like(robot.segment_length), - density=900.0 * jnp.ones_like(robot.segment_length), - damping_coefficient=2.0 * jnp.ones_like(robot.segment_length), - radius_initial=0.022 * jnp.ones_like(robot.segment_length), - radius_final=0.022 * jnp.ones_like(robot.segment_length), + shear_modulus=(1.25e6 / 2.9) * jnp.ones_like(robot.segment_length), + material_damping_coefficient=2.0 * jnp.ones_like(robot.segment_length), ) ) diff --git a/tests/systems/test_pcs.py b/tests/systems/test_pcs.py index d72bb974..f904e144 100644 --- a/tests/systems/test_pcs.py +++ b/tests/systems/test_pcs.py @@ -165,18 +165,18 @@ def test_constant_strain_call(): test_cases = [ ( jnp.zeros((6,)), - params.length[0] / 2, - jnp.eye(4).at[0, 3].set(params.length[0] / 2), + params.link.length[0] / 2, + jnp.eye(4).at[0, 3].set(params.link.length[0] / 2), ), ( jnp.zeros((6,)), - params.length[0], - jnp.eye(4).at[0, 3].set(params.length[0]), + params.link.length[0], + jnp.eye(4).at[0, 3].set(params.link.length[0]), ), ( jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), - params.length[0], - jnp.eye(4).at[0, 3].set(2 * params.length[0]), + params.link.length[0], + jnp.eye(4).at[0, 3].set(2 * params.link.length[0]), ), ] @@ -216,7 +216,7 @@ def test_constant_strain_call(): ) print("[Valid test]\n") - q = jnp.array([jnp.pi / (2 * params.length[0]), 0.0, 0.0, 0.0, 0.0, 0.0]) + q = jnp.array([jnp.pi / (2 * params.link.length[0]), 0.0, 0.0, 0.0, 0.0, 0.0]) qd = jnp.zeros((6,)) u = jnp.ones((6,)) # identity torque for testing print("q = ", q, "qd = ", qd, "u = ", u) @@ -262,11 +262,11 @@ def test_constant_strain_call(): assert not jnp.isnan(E_pot).any(), "Potential energy contains NaN!" E_pot_th = jnp.array( 0.5 - * params.density[0] + * params.link.density[0] * jnp.pi - * params.radius[0] ** 2 + * params.link.cross_section.coefficients[0, 0] ** 2 * jnp.linalg.norm(params.gravity) - * params.length[0] ** 2 + * params.link.length[0] ** 2 ) assert_allclose(E_pot, E_pot_th, rtol=RTOL, atol=ATOL) print("[Valid test]\n") @@ -292,17 +292,22 @@ def test_public_pcs_accessors_and_geometry_helpers() -> None: q = jnp.zeros((int(model.num_active_strains.item()),), dtype=jnp.float64) assert model.is_planar is False - assert_allclose(model.length, jnp.sum(params.length), rtol=RTOL, atol=ATOL) - assert_allclose(model.segment_length, params.length, rtol=RTOL, atol=ATOL) + assert_allclose(model.length, jnp.sum(params.link.length), rtol=RTOL, atol=ATOL) + assert_allclose(model.segment_length, params.link.length, rtol=RTOL, atol=ATOL) - s_second = params.length[0] + 0.25 * params.length[1] + s_second = params.link.length[0] + 0.25 * params.link.length[1] segment_idx, s_local = model.classify_segment(s_second) assert int(segment_idx) == 1 - assert_allclose(s_local, 0.25 * params.length[1], rtol=RTOL, atol=ATOL) + assert_allclose(s_local, 0.25 * params.link.length[1], rtol=RTOL, atol=ATOL) tag, geom = model.cross_section_geometry(q, s_second) assert int(tag) == CrossSectionGeometry.CIRCULAR - assert_allclose(geom, jnp.array([params.radius[1]]), rtol=RTOL, atol=ATOL) + assert_allclose( + geom, + jnp.array([params.link.cross_section.coefficients[1, 0]]), + rtol=RTOL, + atol=ATOL, + ) @pytest.mark.parametrize("num_segments", [1, 2, 3]) @@ -1365,12 +1370,13 @@ def test_cached_constant_matrices_refresh_after_update_params(): strain_selector=jnp.tile(selector_per_segment, 2), ) - updated = model.update_params( - radius=1.1 * model.r, + updated = model.update_link_params( + cross_section=model.params.link.cross_section.replace( + coefficients=1.1 * model.params.link.cross_section.coefficients + ), density=0.9 * model.rho, - young_modulus=1.25 * model.E, - shear_modulus=0.75 * model.G, - damping_matrix=2.0 * model.D_full, + stiffness=1.25 * model.params.link.stiffness, + damping=2.0 * model.params.link.damping, ) segment_ids = jnp.arange(updated.num_segments) expected_M = jax.vmap(updated._compute_local_mass_matrix)(segment_ids) @@ -1514,7 +1520,7 @@ def test_rotational_strain_basis_length_scaling_matches_unscaled_coordinates(): atol=ATOL, ) - updated = scaled.update_params(length=jnp.array([0.2, 0.3])) + updated = scaled.update_link_params(length=jnp.array([0.2, 0.3])) updated_scale = jnp.array( [5.0, 5.0, 5.0, 1.0, 1.0, 1.0, 10 / 3, 10 / 3, 10 / 3, 1.0, 1.0, 1.0] ) diff --git a/tests/systems/test_planar_pcs.py b/tests/systems/test_planar_pcs.py index 2282df67..1623450d 100644 --- a/tests/systems/test_planar_pcs.py +++ b/tests/systems/test_planar_pcs.py @@ -260,7 +260,7 @@ def test_planar_constant_strain_call(): # Test the differential relation: delta_chi ≈ J * delta_q print("Testing differential relation: delta_chi ≈ J * delta_q") delta_q = jnp.array([EPS, -EPS, 2 * EPS]) - chi_plus = robot.forward_kinematics(q=q + delta_q, s=params.length[0]) + chi_plus = robot.forward_kinematics(q=q + delta_q, s=params.link.length[0]) chi_pred = chi + J @ delta_q assert_allclose(chi_plus, chi_pred, rtol=RTOL, atol=ATOL) print("[Valid test]\n") @@ -286,17 +286,22 @@ def test_public_planar_pcs_accessors_geometry_and_chi() -> None: q = jnp.zeros((int(model.num_active_strains.item()),), dtype=jnp.float64) assert model.is_planar is True - assert_allclose(model.length, jnp.sum(params.length), rtol=RTOL, atol=ATOL) - assert_allclose(model.segment_length, params.length, rtol=RTOL, atol=ATOL) + assert_allclose(model.length, jnp.sum(params.link.length), rtol=RTOL, atol=ATOL) + assert_allclose(model.segment_length, params.link.length, rtol=RTOL, atol=ATOL) - s_second = params.length[0] + 0.25 * params.length[1] + s_second = params.link.length[0] + 0.25 * params.link.length[1] segment_idx, s_local = model.classify_segment(s_second) assert int(segment_idx) == 1 - assert_allclose(s_local, 0.25 * params.length[1], rtol=RTOL, atol=ATOL) + assert_allclose(s_local, 0.25 * params.link.length[1], rtol=RTOL, atol=ATOL) tag, geom = model.cross_section_geometry(q, s_second) assert int(tag) == CrossSectionGeometry.CIRCULAR - assert_allclose(geom, jnp.array([params.radius[1]]), rtol=RTOL, atol=ATOL) + assert_allclose( + geom, + jnp.array([params.link.cross_section.coefficients[1, 0]]), + rtol=RTOL, + atol=ATOL, + ) xi = model.strain(q) assert_allclose(model.chi(xi, s_second), model.forward_kinematics(q, s_second)) @@ -1452,12 +1457,13 @@ def test_cached_constant_matrices_refresh_after_update_params_planar(): structure=PlanarPCSStructure(strain_selector=jnp.tile(selector_per_segment, 2)), ) - updated = model.update_params( - radius=1.1 * model.r, + updated = model.update_link_params( + cross_section=model.params.link.cross_section.replace( + coefficients=1.1 * model.params.link.cross_section.coefficients + ), density=0.9 * model.rho, - young_modulus=1.25 * model.E, - shear_modulus=0.75 * model.G, - damping_matrix=2.0 * model.D_full, + stiffness=1.25 * model.params.link.stiffness, + damping=2.0 * model.params.link.damping, ) segment_ids = jnp.arange(updated.num_segments) expected_M = jax.vmap(updated._compute_local_mass_matrix)(segment_ids) @@ -1713,7 +1719,7 @@ def test_rotational_strain_basis_length_scaling_matches_unscaled_coordinates_pla atol=ATOL, ) - updated = scaled.update_params(length=jnp.array([0.2, 0.3])) + updated = scaled.update_link_params(length=jnp.array([0.2, 0.3])) updated_scale = jnp.array([5.0, 1.0, 1.0, 10 / 3, 1.0, 1.0]) assert_allclose( updated.B_xi, diff --git a/tests/systems/test_pressure_actuated_pcs_models.py b/tests/systems/test_pressure_actuated_pcs_models.py index 19da0cc4..e7b6adf3 100644 --- a/tests/systems/test_pressure_actuated_pcs_models.py +++ b/tests/systems/test_pressure_actuated_pcs_models.py @@ -9,25 +9,51 @@ from soromox.actuation import ThreadlikeActuator from soromox.systems import ( + PCS, + ContinuumLinkParams, + CrossSectionParams, ISupport, ISupportParams, ISupportStructure, + LinkSpec, ) def make_isupport_params(num_segments=1): angles = 2.0 * jnp.pi * jnp.arange(3) / 3 + length = jnp.full((num_segments,), 0.1) + radius = jnp.full((num_segments,), 0.02) + young = jnp.full((num_segments,), 2e3) + shear = jnp.full((num_segments,), 1e3) + area = jnp.pi * radius**2 + transverse = jnp.pi * radius**4 / 4.0 + polar = 2.0 * transverse + stiffness = length[:, None, None] * jax.vmap(jnp.diag)( + jnp.stack( + [ + shear * polar, + young * transverse, + young * transverse, + young * area, + shear * area, + shear * area, + ], + axis=1, + ) + ) return ISupportParams( base_pose=spatial_base_pose(), - length=jnp.full((num_segments,), 0.1), - radius=jnp.full((num_segments,), 0.02), - density=jnp.full((num_segments,), 1000.0), gravity=jnp.array([0.0, 0.0, -9.81]), - young_modulus=jnp.full((num_segments,), 2e3), - shear_modulus=jnp.full((num_segments,), 1e3), - damping_matrix=1e-3 * jnp.eye(6 * num_segments), - reference_strain=jnp.tile( - jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), num_segments + link=ContinuumLinkParams( + length=length, + density=jnp.full((num_segments,), 1000.0), + reference_strain=jnp.tile( + jnp.array([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0]]), + (num_segments, 1), + ), + cross_section=CrossSectionParams(coefficients=radius[:, None]), + stiffness=stiffness, + damping=1e-3 * jnp.tile(jnp.eye(6)[None, :, :], (num_segments, 1, 1)), ), chamber_inner_radius=jnp.full((num_segments,), 0.002), chamber_outer_radius=jnp.full((num_segments,), 0.004), @@ -54,7 +80,7 @@ def expected_isupport_segment_actuation( else: chamber_area = params.chamber_effective_pressure_area[pneumatic_segment_idx] if segment_length is None: - segment_length = params.length[pneumatic_segment_idx] + segment_length = params.link.length[pneumatic_segment_idx] chamber_distance = params.chamber_distance[pneumatic_segment_idx] chamber_angles = params.chamber_azimuth_angles[pneumatic_segment_idx] expected_columns = [] @@ -314,8 +340,8 @@ def test_isupport_body_update_preserves_actuator_owned_bounds(): upper_bounds=upper_bounds, ) - updated_body = updated_actuator.update_params( - density=updated_actuator.params.density * 1.01 + updated_body = updated_actuator.update_link_params( + density=updated_actuator.params.link.density * 1.01 ) assert jnp.allclose(updated_body.actuators[0].params.lower_bounds, lower_bounds) @@ -408,13 +434,29 @@ def test_isupport_actuation_matrix_projects_to_active_strains(): def make_mixed_isupport_params(): angles = 2.0 * jnp.pi * jnp.arange(3) / 3 + lengths = [0.01, 0.10, 0.02, 0.20, 0.03] + radii = [0.011, 0.020, 0.012, 0.025, 0.013] + densities = [1200.0, 1000.0, 1300.0, 1050.0, 1400.0] + young = [2e9, 2e3, 2e9, 3e3, 2e9] + shear = [8e8, 1e3, 8e8, 1.5e3, 8e8] + link = PCS.params_from_links( + [ + LinkSpec.circular( + length=length, + radius=radius, + density=density, + young_modulus=young_modulus, + shear_modulus=shear_modulus, + material_damping_coefficient=1e-3, + reference_strain=[0, 0, 0, 1, 0, 0], + ) + for length, radius, density, young_modulus, shear_modulus in zip( + lengths, radii, densities, young, shear, strict=True + ) + ] + ).link return make_isupport_params(num_segments=5).replace( - length=jnp.array([0.01, 0.10, 0.02, 0.20, 0.03]), - radius=jnp.array([0.011, 0.020, 0.012, 0.025, 0.013]), - density=jnp.array([1200.0, 1000.0, 1300.0, 1050.0, 1400.0]), - young_modulus=jnp.array([2e9, 2e3, 2e9, 3e3, 2e9]), - shear_modulus=jnp.array([8e8, 1e3, 8e8, 1.5e3, 8e8]), - damping_matrix=1e-3 * jnp.eye(30), + link=link, chamber_inner_radius=jnp.array([0.002, 0.003]), chamber_outer_radius=jnp.array([0.004, 0.005]), chamber_distance=jnp.array([0.01, 0.015]), @@ -519,9 +561,20 @@ def test_isupport_rigid_geometry_uses_filled_cylinders(): def test_isupport_preserves_rigid_material_damping_properties(): - params = make_mixed_isupport_params().replace( - damping_matrix=None, - material_damping_coefficient=jnp.array([10.0, 20.0, 30.0, 40.0, 50.0]), + params = make_mixed_isupport_params() + damping_coefficients = jnp.array([10.0, 20.0, 30.0, 40.0, 50.0]) + radius = params.link.cross_section.coefficients[:, 0] + area = jnp.pi * radius**2 + transverse = jnp.pi * radius**4 / 4.0 + polar = 2.0 * transverse + damping = params.link.length[:, None, None] * jax.vmap(jnp.diag)( + damping_coefficients[:, None] + * jnp.stack( + [polar, 3 * transverse, 3 * transverse, 3 * area, area, area], axis=1 + ) + ) + params = params.replace( + link=params.link.replace(damping=damping), ) robot = ISupport( params, @@ -532,12 +585,8 @@ def test_isupport_preserves_rigid_material_damping_properties(): ), ) - assert jnp.allclose( - robot.pcs_params.material_damping_coefficient, - jnp.array([10.0, 20.0, 20.0, 30.0, 40.0, 50.0]), - ) - rigid_area = jnp.pi * params.radius[0] ** 2 - expected_axial_damping = params.length[0] * 10.0 * 3.0 * rigid_area + rigid_area = jnp.pi * params.link.cross_section.coefficients[0, 0] ** 2 + expected_axial_damping = params.link.length[0] * 10.0 * 3.0 * rigid_area assert jnp.allclose(robot.D_full[3, 3], expected_axial_damping) assert jnp.allclose(robot.D_active, robot.B_xi.T @ robot.D_full @ robot.B_xi) @@ -606,14 +655,22 @@ def test_isupport_layout_validation_and_updates(): params.replace(pcs_segment_lengths=jnp.array([0.03, 0.06, 0.20])), structure=structure, ) - bad_reference = params.reference_strain.at[0].set(0.1) + bad_reference = params.link.reference_strain.at[0, 0].set(0.1) with pytest.raises(ValueError, match="Rigid segment reference_strain"): - ISupport(params.replace(reference_strain=bad_reference), structure=structure) + ISupport( + params.replace(link=params.link.replace(reference_strain=bad_reference)), + structure=structure, + ) - updated = robot.update_params( + updated_link = params.link.replace( length=jnp.array([0.015, 0.12, 0.025, 0.20, 0.035]), + cross_section=params.link.cross_section.replace( + coefficients=jnp.array([0.015, 0.020, 0.016, 0.025, 0.017])[:, None] + ), + ) + updated = robot.update_params( + link=updated_link, pcs_segment_lengths=jnp.array([0.05, 0.07, 0.20]), - radius=jnp.array([0.015, 0.020, 0.016, 0.025, 0.017]), ) assert updated.structure.rigid_segment_selector == structure.rigid_segment_selector assert jnp.allclose(updated.L, jnp.array([0.015, 0.05, 0.07, 0.025, 0.20, 0.035])) diff --git a/tests/systems/test_shared_continuum_components.py b/tests/systems/test_shared_continuum_components.py new file mode 100644 index 00000000..0d50e032 --- /dev/null +++ b/tests/systems/test_shared_continuum_components.py @@ -0,0 +1,246 @@ +import jax +import jax.numpy as jnp +import pytest +from numpy.testing import assert_allclose + +from soromox.systems import ( + GVS, + PCS, + GVSSegment, + IsotropicMaterialParams, + JointSpec, + LinearProfile, + LinkSpec, + StrainBasisSpec, +) + +jax.config.update("jax_enable_x64", True) + +REFERENCE = [0.0, 0.0, 0.0, 1.0, 0.0, 0.0] + + +def _material() -> IsotropicMaterialParams: + return IsotropicMaterialParams( + young_modulus=jnp.array([1.0e6]), + shear_modulus=jnp.array([3.4e5]), + material_damping_coefficient=jnp.array([1.0e4]), + ) + + +def _pcs() -> PCS: + return PCS.from_links( + [ + LinkSpec.circular( + length=0.2, + radius=0.012, + density=1000.0, + young_modulus=1.0e6, + shear_modulus=3.4e5, + material_damping_coefficient=1.0e4, + reference_strain=REFERENCE, + ) + ] + ) + + +def _gvs() -> GVS: + return GVS.from_segments( + [ + GVSSegment( + link=LinkSpec.rectangular( + length=0.2, + height=LinearProfile(0.03, 0.02), + width=0.025, + density=1000.0, + young_modulus=1.0e6, + shear_modulus=3.4e5, + material_damping_coefficient=1.0e4, + reference_strain=REFERENCE, + ), + joint=JointSpec.revolute( + "z", + stiffness=jnp.array([[0.3]]), + damping=jnp.array([[0.02]]), + ), + basis=StrainBasisSpec( + type="legendre", + strain_selector=("kappa_y", "sigma_x"), + basis_order=1, + ), + num_gauss_points=7, + ) + ] + ) + + +def test_shared_params_replace_and_explicit_matrix_bypass() -> None: + stiffness = jnp.diag(jnp.arange(1.0, 7.0)) + damping = 0.1 * stiffness + robot = PCS.from_links( + [ + LinkSpec.circular( + length=0.2, + radius=0.01, + density=1000.0, + stiffness=stiffness, + damping=damping, + reference_strain=REFERENCE, + ) + ] + ) + assert_allclose(robot.params.link.stiffness[0], stiffness) + assert_allclose(robot.params.link.damping[0], damping) + updated = robot.update_link_params(density=jnp.array([1050.0])) + assert_allclose(updated.params.link.density, jnp.array([1050.0])) + assert_allclose(robot.params.link.density, jnp.array([1000.0])) + + +def test_gvs_joint_and_link_blocks_both_contribute() -> None: + robot = _gvs() + assert_allclose( + robot.K_full[: robot.max_dof, : robot.max_dof], robot.params.joint.stiffness[0] + ) + assert_allclose( + robot.D_full[: robot.max_dof, : robot.max_dof], robot.params.joint.damping[0] + ) + assert_allclose( + robot.K_full[robot.max_dof :, robot.max_dof :], robot.params.link.stiffness[0] + ) + assert_allclose( + robot.D_full[robot.max_dof :, robot.max_dof :], robot.params.link.damping[0] + ) + + +def test_geometry_refreshes_operators_without_overwriting_matrices() -> None: + robot = _pcs() + original_stiffness = robot.params.link.stiffness + original_operator = robot.young_stiffness_operator + section = robot.params.link.cross_section.replace( + coefficients=1.1 * robot.params.link.cross_section.coefficients + ) + geometry_robot = robot.update_link_params(cross_section=section) + assert_allclose(geometry_robot.params.link.stiffness, original_stiffness) + assert not jnp.allclose(geometry_robot.young_stiffness_operator, original_operator) + rebuilt = geometry_robot.with_isotropic_material(_material()) + assert not jnp.allclose(rebuilt.params.link.stiffness, original_stiffness) + + +def test_material_gradients_and_jit_for_pcs_and_gvs() -> None: + for robot in (_pcs(), _gvs()): + material = _material() + + def objective(candidate, robot=robot): + updated = robot.with_isotropic_material(candidate) + return jnp.sum(updated.params.link.stiffness) + jnp.sum( + updated.params.link.damping + ) + + gradient = jax.jit(jax.grad(objective))(material) + assert_allclose( + gradient.young_modulus, + jnp.sum(robot.young_stiffness_operator, axis=(1, 2)), + ) + assert_allclose( + gradient.shear_modulus, + jnp.sum(robot.shear_stiffness_operator, axis=(1, 2)), + ) + assert_allclose( + gradient.material_damping_coefficient, + jnp.sum(robot.material_damping_operator, axis=(1, 2)), + ) + + +def test_scalar_and_per_link_material_inputs() -> None: + robot = PCS.from_links( + [ + LinkSpec.circular( + length=0.2, + radius=0.012, + density=1000.0, + young_modulus=1.0e6, + shear_modulus=3.4e5, + material_damping_coefficient=1.0e4, + reference_strain=REFERENCE, + ), + LinkSpec.circular( + length=0.15, + radius=0.01, + density=1050.0, + young_modulus=1.0e6, + shear_modulus=3.4e5, + material_damping_coefficient=1.0e4, + reference_strain=REFERENCE, + ), + ] + ) + scalar = IsotropicMaterialParams( + young_modulus=jnp.array(1.0e6), + shear_modulus=jnp.array(3.4e5), + material_damping_coefficient=jnp.array(1.0e4), + ) + per_link = IsotropicMaterialParams( + young_modulus=jnp.full((2,), 1.0e6), + shear_modulus=jnp.full((2,), 3.4e5), + material_damping_coefficient=jnp.full((2,), 1.0e4), + ) + scalar_matrices = robot.link_matrices_from_material(scalar) + per_link_matrices = robot.link_matrices_from_material(per_link) + assert_allclose(scalar_matrices[0], per_link_matrices[0]) + assert_allclose(scalar_matrices[1], per_link_matrices[1]) + + with pytest.raises(ValueError, match=r"shape \(2,\)"): + robot.link_matrices_from_material( + per_link.replace(young_modulus=jnp.ones((3,))) + ) + + +def test_material_gradient_matches_finite_difference() -> None: + robot = _pcs() + material = _material() + + def objective(young): + candidate = material.replace(young_modulus=young) + return jnp.sum(robot.link_matrices_from_material(candidate)[0]) + + automatic = jax.grad(objective)(material.young_modulus) + step = 1.0 + finite_difference = ( + objective(material.young_modulus + step) + - objective(material.young_modulus - step) + ) / (2.0 * step) + assert_allclose(automatic[0], finite_difference, rtol=1e-7, atol=1e-10) + + +def test_log_material_optax_loop_reduces_loss() -> None: + import optax + + robot = _pcs() + material = _material() + target = robot.link_matrices_from_material( + material.replace( + young_modulus=1.1 * material.young_modulus, + shear_modulus=0.9 * material.shear_modulus, + material_damping_coefficient=(1.2 * material.material_damping_coefficient), + ) + ) + log_material = jax.tree.map(jnp.log, material) + + def decode(values): + return jax.tree.map(jnp.exp, values) + + def loss(values): + stiffness, damping = robot.link_matrices_from_material(decode(values)) + return jnp.mean(((stiffness - target[0]) / 100.0) ** 2) + jnp.mean( + ((damping - target[1]) / 10.0) ** 2 + ) + + optimizer = optax.adam(1e-2) + state = optimizer.init(log_material) + value_and_grad = jax.jit(jax.value_and_grad(loss)) + initial = loss(log_material) + for _ in range(25): + _, gradient = value_and_grad(log_material) + updates, state = optimizer.update(gradient, state, log_material) + log_material = optax.apply_updates(log_material, updates) + assert jnp.isfinite(loss(log_material)) + assert loss(log_material) < initial diff --git a/tests/systems/test_soft_robot_defaults.py b/tests/systems/test_soft_robot_defaults.py index 9d7f729f..24de5c9e 100644 --- a/tests/systems/test_soft_robot_defaults.py +++ b/tests/systems/test_soft_robot_defaults.py @@ -8,7 +8,8 @@ custom_jvp_mode, set_custom_jvp_enabled, ) -from soromox.systems.soft_robot import CrossSectionGeometry, SoftRobot +from soromox.systems.components import CrossSectionGeometry +from soromox.systems.soft_robot import SoftRobot from soromox.utils.geometry import poses jax.config.update("jax_enable_x64", True) diff --git a/tests/systems/test_system_lengths.py b/tests/systems/test_system_lengths.py index a67fd675..8233d189 100644 --- a/tests/systems/test_system_lengths.py +++ b/tests/systems/test_system_lengths.py @@ -23,18 +23,20 @@ GVS, PCS, ArticulatedSoftRobot, - CrossSectionGeometry, + GVSSegment, ISupport, ISupportParams, ISupportStructure, + JointSpec, + LinkSpec, PCSStructure, Pendulum, PlanarHSA, PlanarHSAParams, PlanarHSAStructure, PlanarPCS, + StrainBasisSpec, ) -from soromox.systems.gvs import GVSSegment, JointSpec, LinkSpec, StrainBasisSpec jax.config.update("jax_enable_x64", True) @@ -104,22 +106,20 @@ def _segments(): lengths = [0.11, 0.13, 0.17] return [ GVSSegment( - link=LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=1e6, - nu=0.45, - rho=1000.0, - eta=0.0, - L=length, - r_i=0.02, - r_f=0.02, + link=LinkSpec.circular( + young_modulus=1e6, + shear_modulus=1e6 / 2.9, + density=1000.0, + material_damping_coefficient=0.0, + length=length, + radius=0.02, + reference_strain=[0, 0, 0, 1, 0, 0], ), joint=JointSpec(type="fixed"), basis=StrainBasisSpec( type="monomial", - active=[1, 1, 1, 1, 1, 1], - orders=[0, 0, 0, 0, 0, 0], - xi_ref=[0, 0, 0, 1, 0, 0], + strain_selector=[1, 1, 1, 1, 1, 1], + basis_order=[0, 0, 0, 0, 0, 0], ), num_gauss_points=5, ) diff --git a/tests/systems/test_typed_params_api.py b/tests/systems/test_typed_params_api.py index 0de2c33a..7192037a 100644 --- a/tests/systems/test_typed_params_api.py +++ b/tests/systems/test_typed_params_api.py @@ -17,19 +17,13 @@ from soromox.actuation import ( ArticulatedTendonImpedance, - ThreadlikeActuator, - ThreadlikeRouting, ) from soromox.systems import ( PCS, ArticulatedSoftRobotParams, - ISupport, - ISupportParams, - ISupportStructure, PCSParams, Pendulum, PendulumParams, - PlanarPCS, PlanarPCSParams, ) from soromox.utils.array_math import blk_diag @@ -266,9 +260,9 @@ def test_params_are_pytrees_and_replace_is_immutable(): assert any(leaf.shape == (2,) for leaf in leaves) - updated = params.replace(length=2.0 * params.length) - assert_allclose(params.length, jnp.array([0.1, 0.1])) - assert_allclose(updated.length, jnp.array([0.2, 0.2])) + updated = params.replace(link=params.link.replace(length=2.0 * params.link.length)) + assert_allclose(params.link.length, jnp.array([0.1, 0.1])) + assert_allclose(updated.link.length, jnp.array([0.2, 0.2])) with pytest.raises(KeyError, match="Unknown parameter field"): params.replace(not_a_field=jnp.array([1.0])) @@ -279,206 +273,33 @@ def test_params_are_pytrees_and_replace_is_immutable(): def test_system_update_rejects_static_shape_changes(): robot = PCS(params=_pcs_params(num_segments=2)) - updated = robot.update_params(length=jnp.array([0.12, 0.13])) + updated = robot.update_link_params(length=jnp.array([0.12, 0.13])) assert_allclose(updated.segment_length, jnp.array([0.12, 0.13])) assert_allclose(robot.segment_length, jnp.array([0.1, 0.1])) - with pytest.raises(ValueError, match="length"): - robot.update_params(length=jnp.array([0.1, 0.1, 0.1])) - with pytest.raises(ValueError, match="radius"): - robot.update_params(radius=jnp.array([0.03])) + with pytest.raises(ValueError, match="shape"): + robot.update_link_params(length=jnp.array([0.1, 0.1, 0.1])) + with pytest.raises(ValueError, match="cross-section|coefficients"): + robot.update_link_params( + cross_section=robot.params.link.cross_section.replace( + coefficients=jnp.array([[0.03]]) + ) + ) with pytest.raises(KeyError, match="Unknown parameter field"): robot.update_params(unknown=jnp.array([0.0])) -def test_pcs_material_damping_coefficient_builds_full_matrix(): - length = jnp.array([0.1, 0.2], dtype=jnp.float64) - radius = jnp.array([0.02, 0.03], dtype=jnp.float64) - coefficient = jnp.array([2.0, 3.0], dtype=jnp.float64) - params = PCSParams( - base_pose=spatial_base_pose(), - length=length, - radius=radius, - density=1000.0 * jnp.ones((2,), dtype=jnp.float64), - young_modulus=1e6 * jnp.ones((2,), dtype=jnp.float64), - shear_modulus=1e5 * jnp.ones((2,), dtype=jnp.float64), - material_damping_coefficient=coefficient, - gravity=jnp.array([0.0, 0.0, -9.81], dtype=jnp.float64), - reference_strain=jnp.tile( - jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0], dtype=jnp.float64), 2 - ), - ) - robot = PCS(params=params) - expected = _expected_spatial_material_damping(length, radius, coefficient) - - assert_allclose(robot.D_full, expected) - assert_allclose(robot.damping_matrix(jnp.zeros((robot.num_dofs,))), expected) - - -def test_planar_pcs_material_damping_coefficient_builds_full_matrix(): - length = jnp.array([0.1, 0.2], dtype=jnp.float64) - radius = jnp.array([0.02, 0.03], dtype=jnp.float64) - coefficient = jnp.array(2.0, dtype=jnp.float64) - params = PlanarPCSParams( - base_pose=planar_base_pose(), - length=length, - radius=radius, - density=1000.0 * jnp.ones((2,), dtype=jnp.float64), - young_modulus=1e6 * jnp.ones((2,), dtype=jnp.float64), - shear_modulus=1e5 * jnp.ones((2,), dtype=jnp.float64), - material_damping_coefficient=coefficient, - gravity=jnp.array([0.0, -9.81], dtype=jnp.float64), - reference_strain=jnp.tile(jnp.array([0.0, 1.0, 0.0], dtype=jnp.float64), 2), - ) - robot = PlanarPCS(params=params) - expected = _expected_planar_material_damping(length, radius, coefficient) - - assert_allclose(robot.D_full, expected) - assert_allclose(robot.damping_matrix(jnp.zeros((robot.num_dofs,))), expected) - - -def test_pcs_damping_input_validation(): - kwargs = { - "base_pose": spatial_base_pose(), - "length": jnp.array([0.1], dtype=jnp.float64), - "radius": jnp.array([0.02], dtype=jnp.float64), - "density": jnp.array([1000.0], dtype=jnp.float64), - "young_modulus": jnp.array([1e6], dtype=jnp.float64), - "shear_modulus": jnp.array([1e5], dtype=jnp.float64), - "gravity": jnp.array([0.0, 0.0, -9.81], dtype=jnp.float64), - "reference_strain": jnp.array( - [0.0, 0.0, 0.0, 1.0, 0.0, 0.0], dtype=jnp.float64 - ), - } - - with pytest.raises(ValueError, match="Exactly one"): - PCSParams(**kwargs).validate() - with pytest.raises(ValueError, match="Exactly one"): - PCSParams( - **kwargs, - damping_matrix=jnp.eye(6, dtype=jnp.float64), - material_damping_coefficient=jnp.array([1.0], dtype=jnp.float64), - ).validate() - with pytest.raises(ValueError, match="material_damping_coefficient"): - PCSParams( - **kwargs, - material_damping_coefficient=jnp.array([1.0, 2.0], dtype=jnp.float64), - ).validate() - with pytest.raises(ValueError, match="damping_matrix"): - PCSParams( - **kwargs, - damping_matrix=jnp.eye(5, dtype=jnp.float64), - ).validate() - - -def test_material_damping_updates_and_matrix_switching(): - params = _pcs_params(num_segments=1).replace( - damping_matrix=None, - material_damping_coefficient=jnp.array([1.0], dtype=jnp.float64), - ) - robot = PCS(params=params) - expected = _expected_spatial_material_damping( - params.length, params.radius, params.material_damping_coefficient - ) - assert_allclose(robot.D_full, expected) - - updated = robot.update_params( - material_damping_coefficient=jnp.array([2.0], dtype=jnp.float64) - ) - assert_allclose(updated.D_full, 2.0 * expected) - - damping_matrix = 0.5 * jnp.eye(6, dtype=jnp.float64) - matrix_updated = updated.update_params( - material_damping_coefficient=None, - damping_matrix=damping_matrix, - ) - assert_allclose(matrix_updated.D_full, damping_matrix) - - coefficient_updated = matrix_updated.update_params( - damping_matrix=None, - material_damping_coefficient=jnp.array([3.0], dtype=jnp.float64), - ) - assert_allclose(coefficient_updated.D_full, 3.0 * expected) - - -def test_existing_damping_matrix_path_is_unchanged(): +def test_canonical_link_damping_updates_without_material_duplication(): params = _pcs_params(num_segments=2) robot = PCS(params=params) - - assert_allclose(robot.D_full, params.damping_matrix) - - -def test_tendon_pcs_inherits_material_damping_path(): - body = _pcs_params(num_segments=1).replace( - damping_matrix=None, - material_damping_coefficient=jnp.array([2.0], dtype=jnp.float64), - ) - routing = ThreadlikeRouting.linear( - intercept=jnp.array([0.0, 0.005, 0.0], dtype=jnp.float64), - end_segment_index=(0,), - ) - robot = PCS(params=body, actuators=ThreadlikeActuator.tendons(routing)) - - expected = _expected_spatial_material_damping( - body.length, body.radius, body.material_damping_coefficient - ) + expected = blk_diag(params.link.damping) assert_allclose(robot.D_full, expected) - -def test_planar_tendon_pcs_inherits_material_damping_path(): - body = _planar_pcs_params(num_segments=1).replace( - damping_matrix=None, - material_damping_coefficient=jnp.array([2.0], dtype=jnp.float64), - ) - routing = ThreadlikeRouting.linear( - intercept=jnp.array([0.0, 0.005, 0.0], dtype=jnp.float64), - end_segment_index=(0,), - ) - robot = PlanarPCS(params=body, actuators=ThreadlikeActuator.tendons(routing)) - - expected = _expected_planar_material_damping( - body.length, body.radius, body.material_damping_coefficient - ) - assert_allclose(robot.D_full, expected) - - -def test_isupport_inherits_material_damping_path(): - params = ISupportParams( - base_pose=spatial_base_pose(), - length=jnp.array([0.1], dtype=jnp.float64), - radius=jnp.array([0.02], dtype=jnp.float64), - density=jnp.array([1000.0], dtype=jnp.float64), - gravity=jnp.array([0.0, 0.0, -9.81], dtype=jnp.float64), - young_modulus=jnp.array([2e3], dtype=jnp.float64), - shear_modulus=jnp.array([1e3], dtype=jnp.float64), - material_damping_coefficient=jnp.array([2.0], dtype=jnp.float64), - reference_strain=jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0], dtype=jnp.float64), - chamber_inner_radius=jnp.array([0.002], dtype=jnp.float64), - chamber_outer_radius=jnp.array([0.004], dtype=jnp.float64), - chamber_distance=jnp.array([0.01], dtype=jnp.float64), - chamber_azimuth_angles=(2.0 * jnp.pi * jnp.arange(3, dtype=jnp.float64) / 3.0)[ - None, : - ], - ) - robot = ISupport( - params=params, - structure=ISupportStructure( - num_gauss_points=1, rigid_segment_selector=(False,) - ), - ) - - I_i = robot._local_second_moment_of_area(jnp.array(0)) - A_i = robot._local_cross_sectional_area(jnp.array(0)) - expected_diag = ( - params.length[0] - * params.material_damping_coefficient[0] - * jnp.array( - [I_i[0], 3.0 * I_i[1], 3.0 * I_i[2], 3.0 * A_i, A_i, A_i], - dtype=jnp.float64, - ) - ) - assert_allclose(robot.D_full, jnp.diag(expected_diag)) + updated = robot.update_link_params(damping=2.0 * params.link.damping) + assert_allclose(updated.D_full, 2.0 * expected) + assert not hasattr(updated.params, "material_damping_coefficient") + assert not hasattr(updated.params, "damping_matrix") def test_planar_pcs_params_validate_base_pose_shape(): diff --git a/tools/benchmarks/_benchmark_common.py b/tools/benchmarks/_benchmark_common.py index 2d95bdfc..469917f9 100644 --- a/tools/benchmarks/_benchmark_common.py +++ b/tools/benchmarks/_benchmark_common.py @@ -15,16 +15,16 @@ PCS, ArticulatedSoftRobot, ArticulatedSoftRobotParams, - CrossSectionGeometry, - PCSParams, + GVSSegment, + JointSpec, + LinkSpec, PCSStructure, Pendulum, PendulumParams, PlanarPCS, - PlanarPCSParams, PlanarPCSStructure, + StrainBasisSpec, ) -from soromox.systems.gvs import GVSSegment, JointSpec, LinkSpec, StrainBasisSpec Array = jax.Array @@ -149,23 +149,21 @@ def _planar_pcs_factory(num_segments: int, gauss_points: int = 5) -> PlanarPCS: lengths = jnp.full((num_segments,), 0.12) radii = jnp.full((num_segments,), 0.015) rho = 1070.0 * jnp.ones((num_segments,)) - diag_entries = ( - jnp.repeat(jnp.array([[1.0, 200.0, 200.0]]), num_segments, axis=0) - * lengths[:, None] - ).reshape(-1) - params = PlanarPCSParams( + return PlanarPCS.from_links( + [ + LinkSpec.circular( + length=float(lengths[index]), + radius=float(radii[index]), + density=float(rho[index]), + young_modulus=4.0e5, + shear_modulus=1.5e5, + material_damping_coefficient=5.0e-4, + reference_strain=[0.0, 1.0, 0.0], + ) + for index in range(num_segments) + ], base_pose=jnp.array([jnp.pi / 2, 0.0, 0.0]), - length=lengths, - radius=radii, - density=rho, gravity=jnp.array([0.0, 9.81]), - young_modulus=4.0e5 * jnp.ones((num_segments,)), - shear_modulus=1.5e5 * jnp.ones((num_segments,)), - damping_matrix=5.0e-4 * jnp.diag(diag_entries), - reference_strain=jnp.tile(jnp.array([0.0, 1.0, 0.0]), num_segments), - ) - return PlanarPCS( - params=params, structure=PlanarPCSStructure(num_gauss_points=gauss_points), ) @@ -197,27 +195,21 @@ def _pcs_factory(num_segments: int, gauss_points: int = 5) -> PCS: lengths = jnp.full((num_segments,), 0.1) radii = jnp.full((num_segments,), 0.02) rho = 1050.0 * jnp.ones((num_segments,)) - diag_entries = ( - jnp.repeat( - jnp.array([[1.0, 1.0, 1.0, 300.0, 300.0, 300.0]]), num_segments, axis=0 - ) - * lengths[:, None] - ).reshape(-1) - params = PCSParams( + return PCS.from_links( + [ + LinkSpec.circular( + length=float(lengths[index]), + radius=float(radii[index]), + density=float(rho[index]), + young_modulus=6.0e5, + shear_modulus=2.5e5, + material_damping_coefficient=5.0e-4, + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + ) + for index in range(num_segments) + ], base_pose=jnp.array([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), - length=lengths, - radius=radii, - density=rho, gravity=jnp.array([0.0, 0.0, -9.81]), - young_modulus=6.0e5 * jnp.ones((num_segments,)), - shear_modulus=2.5e5 * jnp.ones((num_segments,)), - damping_matrix=5.0e-4 * jnp.diag(diag_entries), - reference_strain=jnp.tile( - jnp.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), num_segments - ), - ) - return PCS( - params=params, structure=PCSStructure(num_gauss_points=gauss_points), ) @@ -251,22 +243,20 @@ def _gvs_segment(strain_basis_order: int, gauss_points: int) -> GVSSegment: raise ValueError("GVS strain-basis order must be non-negative.") return GVSSegment( - link=LinkSpec( - cross_section_geometry=CrossSectionGeometry.CIRCULAR, - E=1.0e6, - nu=0.45, - rho=980.0, - eta=2.5e3, - L=0.25, - r_i=0.02, - r_f=0.02, + link=LinkSpec.circular( + young_modulus=1.0e6, + shear_modulus=1.0e6 / 2.9, + density=980.0, + material_damping_coefficient=2.5e3, + length=0.25, + radius=0.02, + reference_strain=[0, 0, 0, 1, 0, 0], ), joint=JointSpec(type="fixed"), basis=StrainBasisSpec( type="legendre", - active=[1, 1, 1, 1, 1, 1], - orders=int(strain_basis_order), - xi_ref=[0, 0, 0, 1, 0, 0], + strain_selector=[1, 1, 1, 1, 1, 1], + basis_order=int(strain_basis_order), ), num_gauss_points=gauss_points, ) From 9ba3b74b60e9f0daebe71aa6024ed3834f1ca8f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20St=C3=B6lzle?= Date: Mon, 3 Aug 2026 21:03:10 -0400 Subject: [PATCH 2/5] Document parameter API migration --- docs/development/changelog.md | 52 ++++++++ docs/user-guide/parameter-api-migration.md | 145 +++++++++++++++++++++ mkdocs.yml | 1 + 3 files changed, 198 insertions(+) create mode 100644 docs/user-guide/parameter-api-migration.md diff --git a/docs/development/changelog.md b/docs/development/changelog.md index e13e34c6..4caec581 100644 --- a/docs/development/changelog.md +++ b/docs/development/changelog.md @@ -9,10 +9,62 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Shared continuum-component APIs under `soromox.systems.components` for link, + joint, cross-section, and isotropic-material parameters and specifications. +- Common `from_links`, `params_from_links`, `update_link_params`, + `link_matrices_from_material`, and `with_isotropic_material` workflows for + PCS and PlanarPCS, with the same material/update APIs on GVS. +- GVS joint stiffness and damping parameters that participate in global matrix + assembly, plus `update_joint_params` for immutable joint-local updates. +- Differentiable unit-response mappings from Young's modulus, shear modulus, + and material damping to canonical generalized link matrices. +- Continuum-component, parameter-update, material-optimization, and + [parameter-API migration](../user-guide/parameter-api-migration.md) + documentation with complete PCS and GVS examples. + ### Changed +- Harmonized PCS, PlanarPCS, and GVS around nested `params.link` ownership, + descriptive public construction names, shared link/joint specifications, and + canonical per-link generalized stiffness and damping matrices. +- Moved shared `LinkSpec` and `JointSpec` functionality out of the GVS package; + GVS now retains only segment, strain-basis, quadrature, and runtime concepts + specific to variable-strain mechanics. +- Simplified cross-section parameters to coefficient arrays populated by link + factories; constant values and `LinearProfile` cover constant and linearly + varying geometry without a profile-parameter class hierarchy. +- Made isotropic material parameters caller-owned construction/optimization + PyTrees instead of duplicating material and generalized-matrix + representations inside system parameters. +- Made geometry updates refresh material unit-response operators without + silently replacing explicitly supplied canonical matrices. + ### Fixed +- Included stored GVS joint stiffness and damping in global stiffness and + damping assembly. +- Migrated examples, benchmarks, and paper case studies to the harmonized PCS + and GVS parameter APIs. + +### Breaking changes + +- Removed flat PCS and PlanarPCS link/material fields. Access link data through + `params.link`, and construct systems with `from_links` or + `params_from_links`. +- Removed the global PCS damping matrix and cross-link damping coupling. Supply + one generalized damping block per link. +- Removed `GVSLinkParams` and GVS-local exports of shared link, joint, and + cross-section specifications. Import shared types from `soromox.systems` or + `soromox.systems.components`. +- Replaced abbreviated GVS construction fields such as `E`, `nu`, `rho`, + `eta`, `L`, `r_i`, and `r_f` with descriptive material, geometry, and profile + arguments. +- Moved `reference_strain` from `StrainBasisSpec` to `LinkSpec`; basis specs now + describe only the selected strains and basis order. +- No compatibility aliases are provided. See the compact + [PCS/GVS parameter migration guide](../user-guide/parameter-api-migration.md) + for direct replacements. + ## [0.2.1] - 2026-08-03 ### Added diff --git a/docs/user-guide/parameter-api-migration.md b/docs/user-guide/parameter-api-migration.md new file mode 100644 index 00000000..d98e2f29 --- /dev/null +++ b/docs/user-guide/parameter-api-migration.md @@ -0,0 +1,145 @@ +# PCS/GVS Parameter Migration + +The harmonized parameter API is a clean break: no compatibility aliases are +provided. PCS, PlanarPCS, and GVS now share link parameters and construction +specifications, while GVS additionally owns joint parameters. + +## Import shared components + +Import reusable link, joint, cross-section, and material types from the common +systems API: + +```python +# Before +from soromox.systems.gvs import JointSpec, LinkSpec + +# After +from soromox.systems import JointSpec, LinkSpec +# Equivalent: from soromox.systems.components import JointSpec, LinkSpec +``` + +`GVSSegment` and `StrainBasisSpec` remain GVS concepts. + +## Construct PCS and PlanarPCS from links + +Replace flat `PCSParams` and `PlanarPCSParams` construction with `from_links` +or `params_from_links`: + +```python +from soromox.systems import PCS, LinkSpec + +robot = PCS.from_links([ + LinkSpec.circular( + length=0.2, + radius=0.012, + density=1000.0, + young_modulus=1.0e6, + shear_modulus=3.4e5, + material_damping_coefficient=1.0e4, + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + ) +]) +``` + +Each link may instead supply explicit generalized `stiffness` and `damping` +matrices. Do not mix an explicit matrix with its material-property source. + +## Update field paths + +Link properties are nested consistently: + +| Before | After | +| --- | --- | +| `params.length` | `params.link.length` | +| `params.density` | `params.link.density` | +| `params.reference_strain` | `params.link.reference_strain` | +| `params.stiffness` | `params.link.stiffness` | +| `params.damping` | `params.link.damping` | + +Use `update_link_params` for common changes or nested `replace` for complete +parameter replacement: + +```python +robot = robot.update_link_params( + damping=0.9 * robot.params.link.damping +) + +robot = robot.with_params( + robot.params.replace( + link=robot.params.link.replace(density=new_density) + ) +) +``` + +PCS damping is now link-local. Split a former block-diagonal global damping +matrix into one block per link; cross-link damping terms are unsupported. + +## Rename GVS construction fields + +| Before | After | +| --- | --- | +| `E` | `young_modulus` | +| `nu` | derive `shear_modulus` with `shear_modulus_from_poisson_ratio` | +| `rho` | `density` | +| `eta` | `material_damping_coefficient` | +| `L` | `length` | +| `r_i`, `r_f` | `radius=LinearProfile(base=..., tip=...)` | + +Move `reference_strain` from `StrainBasisSpec` to `LinkSpec`. Replace legacy +`active` and `orders` basis arguments with `strain_selector` and `basis_order`. + +```python +from soromox.systems import ( + GVSSegment, + JointSpec, + LinearProfile, + LinkSpec, + StrainBasisSpec, +) + +segment = GVSSegment( + link=LinkSpec.circular( + length=0.2, + radius=LinearProfile(base=0.015, tip=0.010), + density=1000.0, + young_modulus=1.0e6, + shear_modulus=3.4e5, + material_damping_coefficient=1.0e4, + reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + ), + joint=JointSpec.fixed(), + basis=StrainBasisSpec( + type="legendre", + strain_selector=("kappa_y", "sigma_x"), + basis_order=1, + ), + num_gauss_points=7, +) +``` + +GVS joint matrices are available at `params.joint.stiffness` and +`params.joint.damping` and now contribute to the assembled dynamics. + +## Update isotropic material values + +Material values are caller-owned optimization variables, not duplicated in +runtime system params: + +```python +from soromox.systems import IsotropicMaterialParams + +material = IsotropicMaterialParams( + young_modulus=[1.0e6], + shear_modulus=[3.4e5], + material_damping_coefficient=[1.0e4], +) +material = material.replace( + young_modulus=1.1 * material.young_modulus +) +robot = robot.with_isotropic_material(material) +``` + +For full construction, update, JAX, Optax, geometry co-optimization, and direct +matrix examples, see [Parameters and Optimization](parameters-and-optimization.md). +For component ownership and validation rules, see +[Continuum Robot Components](../api/systems/continuum-components.md). diff --git a/mkdocs.yml b/mkdocs.yml index 5e853cfa..285aa828 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -62,6 +62,7 @@ nav: - User Guide: - Quick Start: user-guide/quick-start.md - Parameters and Optimization: user-guide/parameters-and-optimization.md + - PCS/GVS Parameter Migration: user-guide/parameter-api-migration.md - Examples: user-guide/examples.md - API Reference: - Overview: api/overview.md From 2de12622ffc6ca28c97dd4d4818404b4c9f64b8c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20St=C3=B6lzle?= Date: Mon, 3 Aug 2026 21:08:00 -0400 Subject: [PATCH 3/5] Keep migration guide out of navigation --- docs/user-guide/parameters-and-optimization.md | 2 ++ mkdocs.yml | 1 - 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/user-guide/parameters-and-optimization.md b/docs/user-guide/parameters-and-optimization.md index b2588f4c..428e3284 100644 --- a/docs/user-guide/parameters-and-optimization.md +++ b/docs/user-guide/parameters-and-optimization.md @@ -8,6 +8,8 @@ interfaces in gradient-based optimization. For the mechanical meaning of links, joints, cross-sections, stiffness, damping, and isotropic materials, see [Continuum Robot Components](../api/systems/continuum-components.md). +If you are updating code written against the previous PCS or GVS API, use the +compact [parameter migration guide](parameter-api-migration.md). ## Choosing the right parameter layer diff --git a/mkdocs.yml b/mkdocs.yml index 285aa828..5e853cfa 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -62,7 +62,6 @@ nav: - User Guide: - Quick Start: user-guide/quick-start.md - Parameters and Optimization: user-guide/parameters-and-optimization.md - - PCS/GVS Parameter Migration: user-guide/parameter-api-migration.md - Examples: user-guide/examples.md - API Reference: - Overview: api/overview.md From f14aca1519e2565ff1ba450a95afad10d4cad9fc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20St=C3=B6lzle?= Date: Mon, 3 Aug 2026 21:09:38 -0400 Subject: [PATCH 4/5] Limit migration guide to changelog --- docs/user-guide/parameters-and-optimization.md | 2 -- 1 file changed, 2 deletions(-) diff --git a/docs/user-guide/parameters-and-optimization.md b/docs/user-guide/parameters-and-optimization.md index 428e3284..b2588f4c 100644 --- a/docs/user-guide/parameters-and-optimization.md +++ b/docs/user-guide/parameters-and-optimization.md @@ -8,8 +8,6 @@ interfaces in gradient-based optimization. For the mechanical meaning of links, joints, cross-sections, stiffness, damping, and isotropic materials, see [Continuum Robot Components](../api/systems/continuum-components.md). -If you are updating code written against the previous PCS or GVS API, use the -compact [parameter migration guide](parameter-api-migration.md). ## Choosing the right parameter layer From 697e329c88aad240d4ab12ce19cafe2da37defb9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20St=C3=B6lzle?= Date: Mon, 3 Aug 2026 22:29:42 -0400 Subject: [PATCH 5/5] fix continuum parameter review findings --- docs/api/systems/continuum-components.md | 11 +- docs/development/changelog.md | 6 + src/soromox/rendering/open3d_renderer.py | 12 +- .../systems/components/cross_sections.py | 44 ++++- src/soromox/systems/components/joints.py | 8 + src/soromox/systems/components/links.py | 82 +++++++++- src/soromox/systems/components/materials.py | 43 ++++- src/soromox/systems/gvs/_assembly.py | 2 +- src/soromox/systems/gvs/construction.py | 44 ++++- src/soromox/systems/gvs/core.py | 7 +- src/soromox/systems/gvs/params.py | 19 ++- src/soromox/systems/soft_robot.py | 4 + tests/systems/test_gvs.py | 2 +- .../test_shared_continuum_components.py | 152 ++++++++++++++++++ tools/benchmarks/_benchmark_common.py | 22 +-- 15 files changed, 419 insertions(+), 39 deletions(-) diff --git a/docs/api/systems/continuum-components.md b/docs/api/systems/continuum-components.md index 3385e75d..8274a080 100644 --- a/docs/api/systems/continuum-components.md +++ b/docs/api/systems/continuum-components.md @@ -38,7 +38,10 @@ link. The static system structure determines how each row is interpreted: The supported solid cross-sections are circular, rectangular, and elliptical. `CrossSectionGeometry` identifies the family, while shared geometry utilities -calculate area and second moments of area. +calculate area and second moments of area. Geometry dimensions are finite and +strictly positive. Shared evaluators and `robot.cross_section_geometry(...)` +use `[height, width]` for rectangular sections and +`[semi_major, semi_minor]` for elliptical sections. `LinkSpec` provides the ergonomic construction interface: @@ -85,7 +88,8 @@ circular, rectangular, and elliptical cross-sections. Stiffness and damping are finite symmetric matrices. They are not required to be diagonal or positive definite, which permits anisotropic, coupled, fitted, -and learned constitutive models. +and learned constitutive models. Link lengths and densities must be finite and +strictly positive; reference strains must be finite. The generalized dimension depends on the system: @@ -118,7 +122,8 @@ material = IsotropicMaterialParams( ``` Each field may be scalar or contain one value per link. Scalar values are -broadcast when the material is applied. +broadcast when the material is applied. Young's and shear moduli must be finite +and strictly positive; material damping must be finite and nonnegative. Geometry, link length, strain basis, rotational scaling, and quadrature are projected into unit-response operators. Material matrices are then evaluated diff --git a/docs/development/changelog.md b/docs/development/changelog.md index 4caec581..75cf371d 100644 --- a/docs/development/changelog.md +++ b/docs/development/changelog.md @@ -43,6 +43,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Included stored GVS joint stiffness and damping in global stiffness and damping assembly. +- Resolved omitted GVS padding values in the advanced `GVS(params, structure)` + constructor and aligned rectangular geometry output with the shared + `[height, width]` cross-section convention. +- Normalized caller-provided isotropic material sequences to JAX arrays and + rejected non-finite or nonphysical link geometry, density, and material + values during concrete construction and replacement. - Migrated examples, benchmarks, and paper case studies to the harmonized PCS and GVS parameter APIs. diff --git a/src/soromox/rendering/open3d_renderer.py b/src/soromox/rendering/open3d_renderer.py index 1819ecbe..d8125c67 100644 --- a/src/soromox/rendering/open3d_renderer.py +++ b/src/soromox/rendering/open3d_renderer.py @@ -367,8 +367,8 @@ def _cross_section_ring_offsets( params = np.asarray(params, dtype=np.float64).reshape(-1) eps = 1e-6 if geom_tag == CrossSectionGeometry.RECTANGULAR: - width = max(float(params[0]) if params.size else 0.0, eps) - height = max(float(params[1]) if params.size > 1 else 0.0, eps) + height = max(float(params[0]) if params.size else 0.0, eps) + width = max(float(params[1]) if params.size > 1 else 0.0, eps) return np.array( [ [0.0, -0.5 * width, -0.5 * height], @@ -937,8 +937,8 @@ def _primitive_from_section( return "cylinder", np.array([radius, radius, 1.0]), True, True return "sphere", np.array([radius, radius, radius]), False, False if geom_tag == CrossSectionGeometry.RECTANGULAR: - width = max(float(params[0]) if params.size else 0.0, eps) - height = max(float(params[1]) if params.size > 1 else 0.0, eps) + height = max(float(params[0]) if params.size else 0.0, eps) + width = max(float(params[1]) if params.size > 1 else 0.0, eps) if mode == "swept": return "box", np.array([width, height, 1.0]), True, True depth = max(width, height) @@ -1153,8 +1153,8 @@ def mat_for(color_rgba: np.ndarray | tuple[float, ...]): elif geom_type == CrossSectionGeometry.RECTANGULAR: if params.size < 2: continue - width = max(float(params[0]), 1e-6) - height = max(float(params[1]), 1e-6) + height = max(float(params[0]), 1e-6) + width = max(float(params[1]), 1e-6) depth = max(width, height) box = _make_box_centered( curve[p], diff --git a/src/soromox/systems/components/cross_sections.py b/src/soromox/systems/components/cross_sections.py index 9fb26bfd..f4a84539 100644 --- a/src/soromox/systems/components/cross_sections.py +++ b/src/soromox/systems/components/cross_sections.py @@ -8,6 +8,7 @@ import jax.numpy as jnp from jax import Array +from jax.errors import ConcretizationTypeError, TracerBoolConversionError from soromox.systems.params import BaseSystemParams @@ -44,6 +45,23 @@ class LinearProfile: base: float tip: float + def __post_init__(self) -> None: + """Validate the endpoint dimensions. + + Raises: + ValueError: If either endpoint is non-finite or not strictly + positive. + """ + values = jnp.asarray([self.base, self.tip]) + try: + valid = bool(jnp.all(jnp.isfinite(values) & (values > 0.0))) + except (ConcretizationTypeError, TracerBoolConversionError): + return + if not valid: + raise ValueError( + "LinearProfile base and tip must be finite and strictly positive." + ) + class CrossSectionParams(BaseSystemParams): """Batched dynamic cross-section coefficients. @@ -62,8 +80,15 @@ class CrossSectionParams(BaseSystemParams): coefficients: Array def __check_init__(self) -> None: + object.__setattr__(self, "coefficients", jnp.asarray(self.coefficients)) self.validate() + def _normalize_replacement(self, name: str, value: object) -> object: + """Normalize replacement coefficient values to JAX arrays.""" + if name == "coefficients": + return jnp.asarray(value) + return value + def validate(self) -> None: """Validate the batched cross-section coefficient array. @@ -72,7 +97,8 @@ def validate(self) -> None: Raises: ValueError: If ``coefficients`` is not a nonempty two-dimensional - array. + array, contains non-finite or negative values, or has a row + without any positive geometric coefficient. """ coefficients = jnp.asarray(self.coefficients) if coefficients.ndim != 2: @@ -82,6 +108,22 @@ def validate(self) -> None: ) if coefficients.shape[0] < 1 or coefficients.shape[1] < 1: raise ValueError("cross-section coefficients must be non-empty.") + try: + finite_and_nonnegative = bool( + jnp.all(jnp.isfinite(coefficients) & (coefficients >= 0.0)) + ) + positive_per_link = bool(jnp.all(jnp.any(coefficients > 0.0, axis=1))) + except (ConcretizationTypeError, TracerBoolConversionError): + return + if not finite_and_nonnegative: + raise ValueError( + "cross-section coefficients must be finite and nonnegative." + ) + if not positive_per_link: + raise ValueError( + "cross-section coefficients must contain a positive dimension " + "for every link." + ) def evaluate_profile( diff --git a/src/soromox/systems/components/joints.py b/src/soromox/systems/components/joints.py index c28de4c6..17e51900 100644 --- a/src/soromox/systems/components/joints.py +++ b/src/soromox/systems/components/joints.py @@ -69,8 +69,16 @@ class JointParams(BaseSystemParams): damping: Array def __check_init__(self) -> None: + object.__setattr__(self, "stiffness", jnp.asarray(self.stiffness)) + object.__setattr__(self, "damping", jnp.asarray(self.damping)) self.validate() + def _normalize_replacement(self, name: str, value: object) -> object: + """Normalize replacement joint-matrix values to JAX arrays.""" + if name in ("stiffness", "damping"): + return jnp.asarray(value) + return value + def validate(self) -> None: """Validate joint matrix shapes, finiteness, and symmetry. diff --git a/src/soromox/systems/components/links.py b/src/soromox/systems/components/links.py index dfff03a1..ed0f2977 100644 --- a/src/soromox/systems/components/links.py +++ b/src/soromox/systems/components/links.py @@ -32,6 +32,22 @@ def _validate_symmetric(name: str, value: Array) -> None: raise ValueError(f"{name} must be symmetric in its trailing dimensions.") +def _validate_finite_domain( + name: str, value: Array | float, *, strictly_positive: bool = False +) -> None: + """Validate concrete numeric values while remaining safe under JAX tracing.""" + array = jnp.asarray(value) + try: + finite = bool(jnp.all(jnp.isfinite(array))) + positive = bool(jnp.all(array > 0.0)) if strictly_positive else True + except (ConcretizationTypeError, TracerBoolConversionError): + return + if not finite: + raise ValueError(f"{name} must contain only finite values.") + if not positive: + raise ValueError(f"{name} must be strictly positive.") + + class ContinuumLinkParams(BaseSystemParams): """Canonical batched dynamic parameters for continuum links. @@ -56,8 +72,16 @@ class ContinuumLinkParams(BaseSystemParams): damping: Array def __check_init__(self) -> None: + for name in ("length", "density", "reference_strain", "stiffness", "damping"): + object.__setattr__(self, name, jnp.asarray(getattr(self, name))) self.validate() + def _normalize_replacement(self, name: str, value: object) -> object: + """Normalize replacement link-array values to JAX arrays.""" + if name in ("length", "density", "reference_strain", "stiffness", "damping"): + return jnp.asarray(value) + return value + def validate(self) -> None: """Validate link-array shapes and canonical matrix properties. @@ -68,7 +92,9 @@ def validate(self) -> None: ValueError: If link fields disagree on ``num_links``, reference strain or coefficient arrays are not two-dimensional, canonical matrices are not equally shaped square batches, or a canonical - matrix contains non-finite values or is not symmetric. + matrix contains non-finite values or is not symmetric, length + or density is non-finite or not strictly positive, or reference + strain contains non-finite values. """ length = jnp.asarray(self.length) if length.ndim != 1 or length.shape[0] < 1: @@ -79,11 +105,18 @@ def validate(self) -> None: raise ValueError( f"density must have shape ({num_links},), got {density.shape}." ) + _validate_finite_domain("length", length, strictly_positive=True) + _validate_finite_domain("density", density, strictly_positive=True) reference_strain = jnp.asarray(self.reference_strain) - if reference_strain.ndim != 2 or reference_strain.shape[0] != num_links: + if ( + reference_strain.ndim != 2 + or reference_strain.shape[0] != num_links + or reference_strain.shape[1] < 1 + ): raise ValueError( "reference_strain must have shape (num_links, strain_dimension)." ) + _validate_finite_domain("reference_strain", reference_strain) self.cross_section.validate() if self.cross_section.coefficients.shape[0] != num_links: raise ValueError("cross-section coefficients must have one row per link.") @@ -156,6 +189,19 @@ class LinkSpec: damping: Array | None = None def __post_init__(self) -> None: + _validate_finite_domain("length", self.length, strictly_positive=True) + _validate_finite_domain("density", self.density, strictly_positive=True) + _validate_finite_domain("reference_strain", self.reference_strain) + reference_strain = jnp.asarray(self.reference_strain) + if reference_strain.ndim != 1 or reference_strain.size < 1: + raise ValueError( + "reference_strain must be a nonempty one-dimensional array." + ) + _validate_finite_domain( + "cross_section_coefficients", + jnp.asarray(self.cross_section_coefficients), + strictly_positive=True, + ) material_stiffness = ( self.young_modulus is not None or self.shear_modulus is not None ) @@ -182,6 +228,26 @@ def __post_init__(self) -> None: raise ValueError( "Provide material_damping_coefficient or an explicit damping matrix." ) + if self.young_modulus is not None: + _validate_finite_domain( + "young_modulus", self.young_modulus, strictly_positive=True + ) + if self.shear_modulus is not None: + _validate_finite_domain( + "shear_modulus", self.shear_modulus, strictly_positive=True + ) + if self.material_damping_coefficient is not None: + _validate_finite_domain( + "material_damping_coefficient", self.material_damping_coefficient + ) + try: + nonnegative_damping = bool( + jnp.all(jnp.asarray(self.material_damping_coefficient) >= 0.0) + ) + except (ConcretizationTypeError, TracerBoolConversionError): + nonnegative_damping = True + if not nonnegative_damping: + raise ValueError("material_damping_coefficient must be nonnegative.") for name in ("stiffness", "damping"): value = getattr(self, name) if value is None: @@ -265,7 +331,9 @@ def circular( Raises: ValueError: If material and explicit sources are missing, incomplete, or supplied together, or if an explicit matrix is - not finite, square, and symmetric. + not finite, square, and symmetric, or if a physical dimension, + density, length, or material scalar is outside its valid + domain. """ return cls._make( geometry=CrossSectionGeometry.CIRCULAR, @@ -317,7 +385,9 @@ def rectangular( Raises: ValueError: If material and explicit sources are missing, incomplete, or supplied together, or if an explicit matrix is - not finite, square, and symmetric. + not finite, square, and symmetric, or if a physical dimension, + density, length, or material scalar is outside its valid + domain. """ return cls._make( geometry=CrossSectionGeometry.RECTANGULAR, @@ -369,7 +439,9 @@ def elliptical( Raises: ValueError: If material and explicit sources are missing, incomplete, or supplied together, or if an explicit matrix is - not finite, square, and symmetric. + not finite, square, and symmetric, or if a physical dimension, + density, length, or material scalar is outside its valid + domain. """ return cls._make( geometry=CrossSectionGeometry.ELLIPTICAL, diff --git a/src/soromox/systems/components/materials.py b/src/soromox/systems/components/materials.py index 3bd70e40..cdf50994 100644 --- a/src/soromox/systems/components/materials.py +++ b/src/soromox/systems/components/materials.py @@ -4,6 +4,7 @@ import jax.numpy as jnp from jax import Array +from jax.errors import ConcretizationTypeError, TracerBoolConversionError from soromox.systems.params import BaseSystemParams @@ -50,16 +51,34 @@ class IsotropicMaterialParams(BaseSystemParams): material_damping_coefficient: Array def __check_init__(self) -> None: + for name in ( + "young_modulus", + "shear_modulus", + "material_damping_coefficient", + ): + object.__setattr__(self, name, jnp.asarray(getattr(self, name))) self.validate() + def _normalize_replacement(self, name: str, value: object) -> object: + """Normalize replacement material values to JAX arrays.""" + if name in ( + "young_modulus", + "shear_modulus", + "material_damping_coefficient", + ): + return jnp.asarray(value) + return value + def validate(self) -> None: - """Validate material field dimensionality. + """Validate material field dimensionality and physical values. Returns: None. Raises: - ValueError: If any material field has more than one dimension. + ValueError: If any material field has more than one dimension; + either modulus is non-finite or not strictly positive; or the + material damping coefficient is non-finite or negative. """ for name in ( "young_modulus", @@ -69,6 +88,26 @@ def validate(self) -> None: value = jnp.asarray(getattr(self, name)) if value.ndim > 1: raise ValueError(f"{name} must be scalar or one-dimensional.") + if value.ndim == 1 and value.shape[0] < 1: + raise ValueError(f"{name} must not be empty.") + try: + finite = bool(jnp.all(jnp.isfinite(value))) + valid_domain = bool( + jnp.all(value >= 0.0) + if name == "material_damping_coefficient" + else jnp.all(value > 0.0) + ) + except (ConcretizationTypeError, TracerBoolConversionError): + continue + if not finite: + raise ValueError(f"{name} must contain only finite values.") + if not valid_domain: + qualifier = ( + "nonnegative" + if name == "material_damping_coefficient" + else "strictly positive" + ) + raise ValueError(f"{name} must be {qualifier}.") def broadcast(self, num_links: int) -> IsotropicMaterialParams: """Broadcast scalar material fields to one value per link. diff --git a/src/soromox/systems/gvs/_assembly.py b/src/soromox/systems/gvs/_assembly.py index a497179c..1695cc63 100644 --- a/src/soromox/systems/gvs/_assembly.py +++ b/src/soromox/systems/gvs/_assembly.py @@ -80,7 +80,7 @@ def assign_gvs_runtime_arrays( ) if max_num_gauss_points < 5: raise ValueError( - f"max_num_gauss_points must be greater than 5, got {max_num_gauss_points}." + f"max_num_gauss_points must be at least 5, got {max_num_gauss_points}." ) _set_model_field(model, "max_num_gauss_points", max_num_gauss_points) diff --git a/src/soromox/systems/gvs/construction.py b/src/soromox/systems/gvs/construction.py index 23fb4ea1..609c7fd7 100644 --- a/src/soromox/systems/gvs/construction.py +++ b/src/soromox/systems/gvs/construction.py @@ -55,6 +55,27 @@ def _link_dof(segment: GVSSegmentStructure) -> int: ) +def _resolve_structure(params: GVSParams, structure: GVSStructure) -> GVSStructure: + """Return a fully resolved GVS structure compatible with ``params``.""" + params.validate_against_structure(structure) + max_dof = ( + int(params.link.stiffness.shape[-1]) + if structure.max_dof is None + else int(structure.max_dof) + ) + max_num_gauss_points = ( + max(segment.num_gauss_points for segment in structure.segments) + if structure.max_num_gauss_points is None + else int(structure.max_num_gauss_points) + ) + return GVSStructure( + segments=structure.segments, + max_dof=max_dof, + max_num_gauss_points=max_num_gauss_points, + scale_rotational_basis_by_length=(structure.scale_rotational_basis_by_length), + ) + + def _basis_at_points( segment: GVSSegmentStructure, points: Array, max_dof: int ) -> Array: @@ -109,7 +130,12 @@ def material_operators_from_params( coefficients, unsupported cross-section data, or invalid quadrature configuration. """ - max_dof = int(structure.max_dof) + params.validate_against_structure(structure) + max_dof = ( + int(params.link.stiffness.shape[-1]) + if structure.max_dof is None + else int(structure.max_dof) + ) young_operators = [] shear_operators = [] damping_operators = [] @@ -264,10 +290,24 @@ def params_and_structure_from_segments( raise ValueError( f"max_dof={layout_max_dof} is smaller than required DOF {required_max_dof}." ) + required_max_num_gauss_points = max( + segment.num_gauss_points for segment in static_segments + ) + if any(segment.num_gauss_points < 5 for segment in static_segments): + raise ValueError("Every GVS segment requires at least 5 Gauss points.") + layout_max_num_gauss_points = ( + required_max_num_gauss_points + if max_num_gauss_points is None + else max_num_gauss_points + ) + if layout_max_num_gauss_points < required_max_num_gauss_points: + raise ValueError( + "max_num_gauss_points must cover every segment quadrature rule." + ) structure = GVSStructure( segments=static_segments, max_dof=layout_max_dof, - max_num_gauss_points=max_num_gauss_points, + max_num_gauss_points=layout_max_num_gauss_points, scale_rotational_basis_by_length=scale_rotational_basis_by_length, ) diff --git a/src/soromox/systems/gvs/core.py b/src/soromox/systems/gvs/core.py index d09d4ee4..decca928 100644 --- a/src/soromox/systems/gvs/core.py +++ b/src/soromox/systems/gvs/core.py @@ -19,6 +19,7 @@ from soromox.systems.gvs._assembly import assign_gvs_runtime_arrays from soromox.systems.gvs._runtime import SegmentRuntimeData from soromox.systems.gvs.construction import ( + _resolve_structure, material_operators_from_params, params_and_structure_from_segments, ) @@ -248,7 +249,7 @@ def __init__( raise TypeError("params must be a GVSParams instance.") if not isinstance(structure, GVSStructure): raise TypeError("structure must be a GVSStructure instance.") - params.validate_against_structure(structure) + structure = _resolve_structure(params, structure) super().__init__(base_pose=params.base_pose, **kwargs) self.params = params self.structure = structure @@ -400,7 +401,7 @@ def cross_section_geometry(self, q: Array, s: Array) -> tuple[Array, Array]: Returns: A tuple ``(geometry, dimensions)``. ``geometry`` is the integer :class:`CrossSectionGeometry` tag. ``dimensions`` contains radius - for a circular section, width and height for a rectangle, or + for a circular section, height and width for a rectangle, or semi-major and semi-minor radii for an ellipse. Raises: @@ -422,7 +423,7 @@ def cross_section_geometry(self, q: Array, s: Array) -> tuple[Array, Array]: height = Link.interpolate_param(x, h_params[0], h_params[1]) width = Link.interpolate_param(x, w_params[0], w_params[1]) tag = jnp.asarray(CrossSectionGeometry.RECTANGULAR, dtype=jnp.int32) - return tag, jnp.array([width, height]) + return tag, jnp.array([height, width]) a_params = self.semi_major_params[segment_idx] b_params = self.semi_minor_params[segment_idx] a_val = Link.interpolate_param(x, a_params[0], a_params[1]) diff --git a/src/soromox/systems/gvs/params.py b/src/soromox/systems/gvs/params.py index c0348744..4d659a55 100644 --- a/src/soromox/systems/gvs/params.py +++ b/src/soromox/systems/gvs/params.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, ClassVar import jax.numpy as jnp +from jax.errors import ConcretizationTypeError, TracerBoolConversionError from soromox.systems.components import ContinuumLinkParams, JointParams from soromox.systems.params import BaseSoftRobotParams, validate_quaternion_base_pose @@ -102,6 +103,18 @@ def validate_against_structure(self, structure: GVSStructure) -> None: f"{expected_cross_section_shape}, got " f"{self.link.cross_section.coefficients.shape}." ) + for index, segment in enumerate(structure.segments): + count = sum(segment.link.cross_section_profile_parameter_counts) + active_coefficients = self.link.cross_section.coefficients[index, :count] + try: + valid_dimensions = bool(jnp.all(active_coefficients > 0.0)) + except (ConcretizationTypeError, TracerBoolConversionError): + continue + if not valid_dimensions: + raise ValueError( + "Active cross-section coefficients must be strictly positive " + f"for GVS segment {index}." + ) joint_dofs = [ Joint.DICT_JOINT_TYPE_DOF[segment.joint.type] @@ -122,7 +135,11 @@ def validate_against_structure(self, structure: GVSStructure) -> None: ) ) required_max_dof = max(joint_dofs + link_dofs) - max_dof = required_max_dof if structure.max_dof is None else structure.max_dof + max_dof = ( + int(self.link.stiffness.shape[-1]) + if structure.max_dof is None + else structure.max_dof + ) if max_dof < required_max_dof: raise ValueError( f"max_dof={max_dof} is smaller than required DOF {required_max_dof}." diff --git a/src/soromox/systems/soft_robot.py b/src/soromox/systems/soft_robot.py index c50166a2..a7775812 100644 --- a/src/soromox/systems/soft_robot.py +++ b/src/soromox/systems/soft_robot.py @@ -262,6 +262,10 @@ def cross_section_geometry(self, q: Array, s: Array) -> tuple[Array, Array]: - CrossSectionGeometry.CIRCULAR - CrossSectionGeometry.RECTANGULAR - CrossSectionGeometry.ELLIPTICAL + + Rectangular parameters use ``[height, width]`` order, matching + :func:`soromox.systems.components.section_properties`. Elliptical + parameters use ``[semi_major, semi_minor]`` order. """ ... diff --git a/tests/systems/test_gvs.py b/tests/systems/test_gvs.py index d6ee1b9a..b5dffd7f 100644 --- a/tests/systems/test_gvs.py +++ b/tests/systems/test_gvs.py @@ -820,7 +820,7 @@ def test_public_gvs_accessors_geometry_and_actuation_matrix() -> None: ) assert int(tag) == CrossSectionGeometry.RECTANGULAR assert_allclose( - geom, jnp.array([expected_width, expected_height]), rtol=RTOL, atol=ATOL + geom, jnp.array([expected_height, expected_width]), rtol=RTOL, atol=ATOL ) s_third = ( diff --git a/tests/systems/test_shared_continuum_components.py b/tests/systems/test_shared_continuum_components.py index 0d50e032..ad3b6a1c 100644 --- a/tests/systems/test_shared_continuum_components.py +++ b/tests/systems/test_shared_continuum_components.py @@ -6,13 +6,16 @@ from soromox.systems import ( GVS, PCS, + CrossSectionGeometry, GVSSegment, + GVSStructure, IsotropicMaterialParams, JointSpec, LinearProfile, LinkSpec, StrainBasisSpec, ) +from soromox.systems.components import section_properties jax.config.update("jax_enable_x64", True) @@ -95,6 +98,119 @@ def test_shared_params_replace_and_explicit_matrix_bypass() -> None: assert_allclose(robot.params.link.density, jnp.array([1000.0])) +def test_material_params_normalize_python_sequences() -> None: + material = IsotropicMaterialParams( + young_modulus=[1.0e6], + shear_modulus=[3.4e5], + material_damping_coefficient=[1.0e4], + ) + + assert isinstance(material.young_modulus, jax.Array) + updated = material.replace(young_modulus=1.1 * material.young_modulus) + assert_allclose(updated.young_modulus, jnp.array([1.1e6])) + + +@pytest.mark.parametrize( + ("field", "value", "match"), + [ + ("length", 0.0, "length must be strictly positive"), + ("length", jnp.nan, "length must contain only finite"), + ("density", -1.0, "density must be strictly positive"), + ("young_modulus", 0.0, "young_modulus must be strictly positive"), + ("shear_modulus", jnp.inf, "shear_modulus must contain only finite"), + ( + "material_damping_coefficient", + -1.0, + "material_damping_coefficient must be nonnegative", + ), + ], +) +def test_link_spec_rejects_nonphysical_values(field, value, match) -> None: + kwargs = { + "length": 0.2, + "radius": 0.012, + "density": 1000.0, + "young_modulus": 1.0e6, + "shear_modulus": 3.4e5, + "material_damping_coefficient": 1.0e4, + "reference_strain": REFERENCE, + } + kwargs[field] = value + + with pytest.raises(ValueError, match=match): + LinkSpec.circular(**kwargs) + + +@pytest.mark.parametrize("radius", [0.0, -0.01, jnp.nan]) +def test_link_spec_rejects_invalid_cross_section_dimensions(radius) -> None: + with pytest.raises(ValueError, match="cross_section_coefficients"): + LinkSpec.circular( + length=0.2, + radius=radius, + density=1000.0, + young_modulus=1.0e6, + shear_modulus=3.4e5, + material_damping_coefficient=1.0e4, + reference_strain=REFERENCE, + ) + + +@pytest.mark.parametrize("endpoints", [(0.0, 0.01), (0.01, -0.01), (jnp.nan, 0.01)]) +def test_linear_profile_rejects_invalid_endpoints(endpoints) -> None: + with pytest.raises(ValueError, match="finite and strictly positive"): + LinearProfile(*endpoints) + + +def test_runtime_link_params_reject_nonphysical_geometry_updates() -> None: + link = _pcs().params.link + + with pytest.raises(ValueError, match="length must be strictly positive"): + link.replace(length=jnp.array([0.0])) + with pytest.raises(ValueError, match="density must contain only finite"): + link.replace(density=jnp.array([jnp.nan])) + with pytest.raises(ValueError, match="finite and nonnegative"): + link.replace( + cross_section=link.cross_section.replace(coefficients=jnp.array([[-0.01]])) + ) + with pytest.raises(ValueError, match="reference_strain.*finite"): + link.replace(reference_strain=link.reference_strain.at[0, 0].set(jnp.nan)) + + +def test_gvs_rejects_zero_active_cross_section_dimension() -> None: + robot = _gvs() + coefficients = robot.params.link.cross_section.coefficients.at[0, 2].set(0.0) + invalid_params = robot.params.replace( + link=robot.params.link.replace( + cross_section=robot.params.link.cross_section.replace( + coefficients=coefficients + ) + ) + ) + + with pytest.raises(ValueError, match="strictly positive.*segment 0"): + GVS(invalid_params, robot.structure) + + +@pytest.mark.parametrize( + ("field", "value", "match"), + [ + ("young_modulus", [0.0], "strictly positive"), + ("shear_modulus", [jnp.nan], "finite"), + ("material_damping_coefficient", [-1.0], "nonnegative"), + ], +) +def test_material_params_reject_nonphysical_values(field, value, match) -> None: + kwargs = { + "young_modulus": [1.0e6], + "shear_modulus": [3.4e5], + "material_damping_coefficient": [1.0e4], + } + kwargs[field] = value + + with pytest.raises(ValueError, match=match): + IsotropicMaterialParams(**kwargs) + + def test_gvs_joint_and_link_blocks_both_contribute() -> None: robot = _gvs() assert_allclose( @@ -111,6 +227,42 @@ def test_gvs_joint_and_link_blocks_both_contribute() -> None: ) +def test_gvs_constructor_resolves_optional_structure_padding() -> None: + source = _gvs() + unresolved = GVSStructure( + segments=source.structure.segments, + max_dof=None, + max_num_gauss_points=None, + scale_rotational_basis_by_length=( + source.structure.scale_rotational_basis_by_length + ), + ) + + rebuilt = GVS(source.params, unresolved) + + assert rebuilt.max_dof == source.max_dof + assert rebuilt.max_num_gauss_points == source.max_num_gauss_points + assert rebuilt.structure.max_dof == source.max_dof + assert rebuilt.structure.max_num_gauss_points == source.max_num_gauss_points + + +def test_gvs_rectangular_geometry_uses_shared_dimension_order() -> None: + robot = _gvs() + q = jnp.zeros(robot.num_dofs) + tag, dimensions = robot.cross_section_geometry(q, 0.25 * robot.segment_length[0]) + expected_height = 0.03 + 0.25 * (0.02 - 0.03) + + assert int(tag) == CrossSectionGeometry.RECTANGULAR + assert_allclose(dimensions, jnp.array([expected_height, 0.025])) + ix, iy, iz, area = section_properties(tag, dimensions) + expected_iy = expected_height * 0.025**3 / 12.0 + expected_iz = 0.025 * expected_height**3 / 12.0 + assert_allclose(area, expected_height * 0.025) + assert_allclose(iy, expected_iy) + assert_allclose(iz, expected_iz) + assert_allclose(ix, expected_iy + expected_iz) + + def test_geometry_refreshes_operators_without_overwriting_matrices() -> None: robot = _pcs() original_stiffness = robot.params.link.stiffness diff --git a/tools/benchmarks/_benchmark_common.py b/tools/benchmarks/_benchmark_common.py index 469917f9..98f625c9 100644 --- a/tools/benchmarks/_benchmark_common.py +++ b/tools/benchmarks/_benchmark_common.py @@ -146,21 +146,18 @@ def _articulated_soft_robot_context( def _planar_pcs_factory(num_segments: int, gauss_points: int = 5) -> PlanarPCS: - lengths = jnp.full((num_segments,), 0.12) - radii = jnp.full((num_segments,), 0.015) - rho = 1070.0 * jnp.ones((num_segments,)) return PlanarPCS.from_links( [ LinkSpec.circular( - length=float(lengths[index]), - radius=float(radii[index]), - density=float(rho[index]), + length=0.12, + radius=0.015, + density=1070.0, young_modulus=4.0e5, shear_modulus=1.5e5, material_damping_coefficient=5.0e-4, reference_strain=[0.0, 1.0, 0.0], ) - for index in range(num_segments) + for _ in range(num_segments) ], base_pose=jnp.array([jnp.pi / 2, 0.0, 0.0]), gravity=jnp.array([0.0, 9.81]), @@ -192,21 +189,18 @@ def _forward_kinematics_at_s(s): def _pcs_factory(num_segments: int, gauss_points: int = 5) -> PCS: - lengths = jnp.full((num_segments,), 0.1) - radii = jnp.full((num_segments,), 0.02) - rho = 1050.0 * jnp.ones((num_segments,)) return PCS.from_links( [ LinkSpec.circular( - length=float(lengths[index]), - radius=float(radii[index]), - density=float(rho[index]), + length=0.1, + radius=0.02, + density=1050.0, young_modulus=6.0e5, shear_modulus=2.5e5, material_damping_coefficient=5.0e-4, reference_strain=[0.0, 0.0, 0.0, 1.0, 0.0, 0.0], ) - for index in range(num_segments) + for _ in range(num_segments) ], base_pose=jnp.array([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), gravity=jnp.array([0.0, 0.0, -9.81]),