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.

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.

Every phase in the manifest is predicted, or the arbitrary stages given instead.

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 and predicted_meshes (paths).

Displacement interpretation

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

Bases: PhysioTwin4DBase

Reconstruct geometry from displacement predictions.

The displacements are added to the subject’s fitted reference mesh — the manifest’s fitted_reference_mesh, or the fitted_reference_mesh argument of the single-subject methods — which keeps the result in that mesh’s world frame. That mesh is what physiotwin4d.WorkflowFitStatisticalModelToPatient produces: PCA shape parameters and a deformable registration to the patient. A surface reconstructed from the shape parameters alone is not a substitute and is not accepted.

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.

Every phase in the manifest is predicted, or the arbitrary stages given instead. The displacements are added to the manifest’s fitted_reference_mesh points. Scoring the result against a ground truth belongs to physiotwin4d.WorkflowEvaluateMovement.

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 and predicted_surfaces (paths).

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

Predict one subject at one stage without a manifest.

The prediction stays in the fitted reference mesh’s world frame, since that is where the displacements are applied.

Parameters:
Return type:

dict[str, Any]

Returns:

Dict with predicted_surface (path) and predicted_points.

process_time_series(shape_parameters, stages, output_directory, fitted_reference_mesh, 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.

  • fitted_reference_mesh (Path) – The subject’s fitted reference mesh, as produced by physiotwin4d.WorkflowFitStatisticalModelToPatient. The displacements are added to its points and the result stays in its world frame.

  • 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 and stage_meshes. Entries that were not requested are empty lists or None.

stage_meshes are the very objects written to predicted_surfaces, handed back so a caller that scores them — physiotwin4d.WorkflowEvaluateMovement — can annotate and re-save them without re-reading. They are retained for the USD writer regardless, so returning them costs nothing.

Raises:

ValueError – If stages is empty.

create_deformation_field(shape_parameters, stage, reference_image, fitted_reference_mesh, output_directory=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 fitted_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. 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).

  • fitted_reference_mesh (Path) – The subject’s fitted reference mesh, as produced by physiotwin4d.WorkflowFitStatisticalModelToPatient. Its points supply the binning positions and normals, and must share the template topology (same point count and ordering).

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

  • 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. WorkflowInferMovement.predict_single() and create_deformation_field() both require a fitted_reference_mesh — the patient’s shape-model surface as fitted by WorkflowFitStatisticalModelToPatient, which is shape parameters and a deformable registration. The prediction stays in that mesh’s world frame, so a fit that carried a pose transform lands where the patient actually is. A surface reconstructed from the shape parameters alone is not a substitute and is not accepted.

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