Predicting With a Mesh-Stage Model

Inference is split in two so that the generic half stays target-agnostic:

  • WorkflowInferPhysicsNeMo loads the checkpoint and returns the raw (n_points, n_target) prediction, whatever the target means.

  • WorkflowInferMovement wraps it and interprets three-component targets as displacements — deformed meshes, error statistics in millimetres, and rasterized deformation fields.

Generic prediction

class physiotwin4d.WorkflowInferPhysicsNeMo(model_directory, inference_method=None, epoch=None, log_level=20)[source]

Bases: PhysioTwin4DBase

Predict per-point targets for a subject at requested stages.

The network is supplied as an inference method — pass a physiotwin4d.InferPhysicsNeMoMGN or physiotwin4d.InferPhysicsNeMoMLP instance as inference_method; a default MeshGraphNet method is used when none is given.

__init__(model_directory, inference_method=None, epoch=None, log_level=20)[source]

Load a trained model and its normalization statistics.

Parameters:
Raises:

  • FileNotFoundError – If the model checkpoint or the template mesh is missing.

  • TypeError – If inference_method is neither None nor an InferPhysicsNeMoBase instance.

property template_mesh: pyvista.DataSet

The shared PCA template mesh defining node order and topology.

load_pca_assets()[source]

Load (and cache) the PCA template mesh and model for reconstruction.

Return type:

tuple[DataSet, dict]

reference_points_from_coefficients(pca_coeffs)[source]

Reconstruct a subject’s reference points in the template’s domain.

The PCA model may be volumetric while the model was trained on the template’s surface (use_template_surface), so the reconstruction is surface-extracted when its point count does not match the template’s.

Return type:

ndarray

predict(pca_coeffs, stage)[source]

Predict (n_points, n_target) targets for a subject at a stage.

Return type:

ndarray

predicted_mesh(targets)[source]

Return a template copy carrying targets as its target array.

Return type:

DataSet

process(subject_manifest, stages=None, output_directory=None)[source]

Predict a subject’s targets from a manifest.

When stages is None every phase in the manifest is predicted and, because the stored target array is available, per-phase error statistics are computed and written. When stages is given those arbitrary stages are predicted without comparison.

Parameters:

  • subject_manifest (Path) – Path to the subject manifest JSON.

  • stages (Optional[list[float]]) – Optional list of stages to predict.

  • output_directory (Optional[Path]) – Output directory; defaults to <model_directory>/<subject_id>.

Return type:

dict[str, Any]

Returns:

Dict with subject_id, predicted_meshes (paths) and, in the phase mode, statistics and statistics_file.

Displacement interpretation

class physiotwin4d.WorkflowInferMovement(inference_workflow, log_level=20)[source]

Bases: PhysioTwin4DBase

Reconstruct geometry from displacement predictions.

The displacements are added to the caller’s reference mesh when one is available — the manifest’s reference_mesh, or the reference_mesh argument of the single-subject methods — which keeps the result in that mesh’s world frame. With no reference mesh they are added to the mesh reconstructed from the PCA coefficients alone, which stays in the model’s PCA frame.

Parameters:
Raises:

ValueError – If the wrapped model does not predict exactly three components, which cannot be a displacement.

__init__(inference_workflow, log_level=20)[source]

Initialize the base class with logging configuration.

Parameters:

  • class_name – Name for the class (used in log messages). If None, uses the class name. Default: None

  • log_level (int | str) – Logging level. Can be an integer (logging.DEBUG, logging.INFO, logging.WARNING, logging.ERROR, logging.CRITICAL) or a string (‘DEBUG’, ‘INFO’, ‘WARNING’, ‘ERROR’, ‘CRITICAL’). Default: logging.INFO

  • log_to_file – Optional file path to write logs to in addition to console output. Default: None

process(subject_manifest, stages=None, output_directory=None)[source]

Predict a subject’s deformed meshes from a manifest.

When stages is None every phase in the manifest is predicted and compared against its ground truth (reference + stored displacement); when stages is given those arbitrary stages are predicted without comparison. The displacements are added to the manifest’s reference_mesh points.

Parameters:

  • subject_manifest (Path) – Path to the subject manifest JSON.

  • stages (Optional[list[float]]) – Optional list of stages to predict.

  • output_directory (Optional[Path]) – Output directory; defaults to <model_directory>/<subject_id>.

Return type:

dict[str, Any]

Returns:

Dict with subject_id, predicted_surfaces (paths), and, in the phase mode, statistics, statistics_file and rmse_surface.

predict_single(shape_parameters, stage, reference_mesh=None, ground_truth=None, output_directory=None)[source]

Predict one subject at one stage without a manifest.

Without reference_mesh the subject reference is reconstructed from the PCA shape parameters (P = mean + Σ b_i·std_i·eigenvector_i) in the SSM/PCA frame, so the prediction is self-consistent with no reference mesh file at all; with one, the prediction stays in that mesh’s world frame.

Parameters:

  • shape_parameters (Path) – JSON file with the subject PCA coefficient vector.

  • stage (float) – Target stage to predict.

  • reference_mesh (Optional[Path]) – The subject’s reference mesh; omit to displace the PCA reconstruction instead.

  • ground_truth (Optional[Path]) – Optional mesh whose points are the true stage positions, for error reporting.

  • output_directory (Optional[Path]) – Output directory; defaults to <model_directory>/single_prediction.

Return type:

dict[str, Any]

Returns:

Dict with predicted_surface (path), predicted_points, and, when ground_truth is supplied, statistics.

process_time_series(shape_parameters, stages, output_directory, reference_mesh=None, ground_truth=None, reference_image=None, warp_interpolation='linear', warp_background_value=0.0, smoothing_sigma_mm=10.0, usd_project_name=None, anatomy_type=None, separate_by_connectivity=False)[source]

Predict one subject across a whole time series and write its geometry.

One prediction per entry of stages, each written as a mesh. When reference_image is supplied, each stage also gets a deformation field, which is Gaussian-smoothed into a continuous itk.DisplacementFieldTransform and used to carry reference_image into that stage’s frame. The smoothing spreads a surface-shell field into the volume, so the warped image is an interpolation of the surface motion, not an independent registration.

Parameters:

  • shape_parameters (Path) – JSON file with the subject PCA coefficient vector.

  • stages (Sequence[float]) – Stages to predict, in the order they are to be animated.

  • output_directory (Path) – Directory every artifact is written to.

  • reference_mesh (Optional[Path]) – The subject’s reference mesh; omit to displace the PCA reconstruction instead.

  • ground_truth (Optional[Sequence[Path]]) – One mesh per stage whose points are the true stage positions, for error reporting. Must align with stages.

  • reference_image (Optional[Image]) – Image carried through each stage’s deformation, and the grid the deformation field is rasterized on. Omit to write meshes only.

  • warp_interpolation (str) – Interpolation used to resample reference_image: "linear" for intensity images, "nearest" for labelmaps and masks.

  • warp_background_value (float) – Value written where a stage’s grid samples outside reference_image. 0.0 suits labelmaps; CT needs -1000.0, which is air in Hounsfield units.

  • smoothing_sigma_mm (float) – Gaussian sigma, in millimeters, that turns the sparse surface-shell field into a continuous deformation.

  • usd_project_name (Optional[str]) – When given, the stage meshes are also written as one animated USD under this name, one time sample per stage.

  • anatomy_type (Optional[str]) – Anatomy whose materials color that USD.

  • separate_by_connectivity (bool) – Whether that USD splits each frame into separate objects by connectivity.

Return type:

dict[str, Any]

Returns:

Dict with stages, predicted_surfaces, warped_images, transforms, usd_file, statistics and statistics_file. Entries that were not requested are empty lists or None.

Raises:

ValueError – If stages is empty, or ground_truth is given with a different length.

create_deformation_field(shape_parameters, stage, reference_image, output_directory=None, reference_mesh=None, direction='forward')[source]

Rasterize the inferred deformation onto a reference image grid.

Each mesh vertex is binned by its reference (undeformed) position into reference_image’s voxel grid. Each voxel of the deformation field holds the mean network displacement (dx, dy, dz) of the vertices that fall in it; each voxel of the normal image holds the mean (renormalized) reference-surface normal of those vertices. Empty voxels are zero.

That is the "forward" field, which maps reference positions to stage positions and is what transforming a mesh needs. Resampling an image, though, maps each output point through the transform to find where to sample the input, so carrying the reference image into the stage frame needs the opposite mapping. direction="inverse" builds it exactly rather than by negating the forward field: each vertex is binned by its deformed position reference + displacement and contributes -displacement.

The binning positions come from reference_mesh, so a patient scan whose statistical-model fit applied a pose transform not captured by the shape coefficients is binned where it actually aligns with reference_image. Omit it to bin at the PCA reconstruction instead, in the model’s own frame. The network displacements themselves depend only on the coefficients and stage, not on the binning positions.

Parameters:

  • shape_parameters (Path) – JSON file with the subject PCA coefficient vector.

  • stage (float) – Target stage for the deformation.

  • reference_image (Image) – The frame’s image; defines the output grid geometry (size, spacing, origin, direction).

  • output_directory (Optional[Path]) – If given, the three images are written there as compressed .mha files.

  • reference_mesh (Optional[Path]) – Mesh whose points supply the binning positions and normals; omit to use the PCA reconstruction. Must share the template topology (same point count and ordering).

  • direction (Literal['forward', 'inverse']) – "forward" for the reference-to-stage field that deforms meshes, "inverse" for the stage-to-reference field that resamples images into the stage frame.

Return type:

dict[str, Any]

Returns:

Dict with deformation_field and normal_image (ITK vector images), weight_image (the vertex count per voxel, which distinguishes an empty voxel from one whose displacement happens to be zero and is what TransformTools.smooth_deformation_field_transform() normalizes by), deformed_surface (the stage mesh as pv.DataSet) and, when written, their paths.

Inference methods

class physiotwin4d.InferPhysicsNeMoBase(log_level=20)[source]

Bases: PhysioTwin4DBase

Base class for a PhysicsNeMo mesh-stage inference method.

Not instantiated directly — use physiotwin4d.InferPhysicsNeMoMGN or physiotwin4d.InferPhysicsNeMoMLP. Subclasses implement build_model(), load_artifacts() and predict(), and set the class attribute model_tag.

model_tag: str = 'base'
__init__(log_level=20)[source]

Initialize the inference method.

Parameters:

log_level (int | str) – Logging level. Default: logging.INFO.

build_model(meta)[source]

Rebuild the (uncompiled) network from checkpoint metadata.

Return type:

Module

load_artifacts(model_directory, n_points, device)[source]

Load any architecture-specific artifacts (MGN graph tensors).

Return type:

None

predict(node_feats)[source]

Run the network over all nodes; return the (n, n_target) output.

Return type:

ndarray

set_model(model, device)[source]

Attach the loaded model and its device before predicting.

Return type:

None

class physiotwin4d.InferPhysicsNeMoMGN(log_level=20)[source]

Bases: InferPhysicsNeMoBase

Predict mesh stages with a trained PhysicsNeMo MeshGraphNet.

The shared graph topology and edge features are loaded from the tensors the training workflow saved next to the checkpoint.

model_tag: str = 'mgn'
build_model(meta)[source]

Rebuild the (uncompiled) network from checkpoint metadata.

Return type:

Module

load_artifacts(model_directory, n_points, device)[source]

Load any architecture-specific artifacts (MGN graph tensors).

Return type:

None

predict(node_feats)[source]

Run the network over all nodes; return the (n, n_target) output.

Return type:

ndarray

class physiotwin4d.InferPhysicsNeMoMLP(log_level=20)[source]

Bases: InferPhysicsNeMoBase

Predict mesh stages with a trained PhysicsNeMo FullyConnected model.

model_tag: str = 'mlp'
build_model(meta)[source]

Rebuild the (uncompiled) network from checkpoint metadata.

Return type:

Module

load_artifacts(model_directory, n_points, device)[source]

Load any architecture-specific artifacts (MGN graph tensors).

Return type:

None

predict(node_feats)[source]

Run the network over all nodes; return the (n, n_target) output.

Return type:

ndarray

Example

from physiotwin4d import (
    InferPhysicsNeMoMGN,
    WorkflowInferMovement,
    WorkflowInferPhysicsNeMo,
)

infer = WorkflowInferPhysicsNeMo(
    model_directory=model_dir,
    inference_method=InferPhysicsNeMoMGN(),   # the default
)

# Raw targets, whatever the model was trained to predict.
targets = infer.predict(pca_coefficients, stage=0.5)

# Or, for a displacement model, geometry and mm error statistics.
movement = WorkflowInferMovement(infer)
result = movement.process(subject_manifest, output_directory=out_dir)

Notes

Where the displacements are applied. Give WorkflowInferMovement.predict_single() or create_deformation_field() a reference_mesh and the prediction stays in that mesh’s world frame — the right choice for a patient scan whose fit carried a pose transform. Omit it and the displacements are applied to the mesh reconstructed from the PCA coefficients alone, which needs no per-subject geometry but stays in the model’s own frame.

Arbitrary stages. Nothing constrains stage to a phase that was acquired; predicting between acquired phases is the reason to train a surrogate at all.

Deformation fields. create_deformation_field() bins the per-vertex displacements and reference-surface normals onto a caller-supplied image grid, giving an ITK vector image you can apply to volumes and labelmaps with TransformTools.

See Also