API reference

Public symbols are defined on the stressnet package. Heavy runtime dependencies (TensorFlow, Spektral, ForSys) may be mocked when building these docs without a full scientific stack.

Top-level package

StressNET: deep learning for mechanical stress inference in tissues.

This package exposes graph construction from ForSys/Surface Evolver inputs, registered model weights (load_model()), inference helpers, and plotting utilities. See the documentation for API details.

Author

Nicolás Aldecoa Rodrigo.

Preprocessing

stressnet.se_output_to_graph(src_file: str | Path, *, include_targets: bool = True, edge_n_vertices: int = 9, apply_savgol_filter: bool = False, include_forsys_predictions: bool = True, forsys_solve_method: str = 'default', render_plots: bool = False, plots_dir: str | Path = './debug', return_timers: bool = True, tag_cell_interfaces: list | None = None, jitter_kwargs: dict | None = None, raise_if_gt_is_zero: bool = True, debug_plots_prefix: str | None = None, return_frame: bool = False, verbose: bool = True) → dict[str, Any][source]

Load a Surface Evolver dump, build a ForSys frame, and return graph arrays.

Parameters:

  • src_file – Path to the Surface Evolver output file.

  • include_targets – Whether ground-truth tensions are present and should populate gt on the frame.

  • edge_n_vertices – Number of vertices sampled along each lattice edge for feature extraction.

  • apply_savgol_filter – If True, smooth resampled vertex coordinates with a Savitzky–Golay filter.

  • include_forsys_predictions – Whether to run ForSys inference for auxiliary predictions in the output dict.

  • forsys_solve_method – Solver label passed to ForSys tension recovery.

  • render_plots – If True, write debug plots under plots_dir.

  • plots_dir – Directory for optional debug figures.

  • return_timers – Include timing fields in the returned dictionary.

  • tag_cell_interfaces – Optional list controlling interface tagging between cells.

  • jitter_kwargs – Optional dict of jitter parameters forwarded to vertex jittering.

  • raise_if_gt_is_zero – If True, validate non-degenerate ground-truth tensions when present.

  • debug_plots_prefix – Filename prefix for plots; defaults to the stem of src_file.

  • return_frame – If True, include the constructed frame object under key 'frame'.

  • verbose – Enable logging of major steps.

Returns:

Keys typically include 'a', 'x', 'e', optional 'y', timing keys, and optional ForSys prediction arrays depending on flags.

Return type:

dict

stressnet.skeleton_to_graph(src_file: str | Path, gt_file: str | Path | None = None, *, mirror_y: bool = False, edge_n_vertices: int = 9, fixed_ne: int | None = None, apply_savgol_filter: bool = False, include_forsys_predictions: bool = True, forsys_solve_method: str = 'default', render_plots: bool = False, plots_dir: str | Path = './debug', return_timers: bool = True, tag_cell_interfaces: list | None = None, raise_if_gt_is_zero: bool = True, debug_plots_prefix: str | None = None, return_frame: bool = False, verbose: bool = True) → dict[str, Any][source]

Load a skeleton .tif, build a mesh in ForSys, optionally attach myosin GT, and return graph tensors.

Parameters:

  • src_file – Path to a binary skeleton image (.tif).

  • gt_file – Optional myosin intensity image (.tif) used to assign ground-truth tensions.

  • mirror_y – Passed to ForSys skeleton loading (flip image rows).

  • edge_n_vertices – Target number of vertices per edge for mesh generation.

  • fixed_ne – If set, overrides derived ne for generate_mesh.

  • apply_savgol_filter – Whether to smooth resampled vertices with Savitzky–Golay filtering.

  • include_forsys_predictions – Whether to include ForSys baseline predictions in the output.

  • forsys_solve_method – ForSys solver label for tension recovery.

  • render_plots – Write debug plots when True.

  • plots_dir – Output directory for debug plots.

  • return_timers – Attach timing metadata to the result dict.

  • tag_cell_interfaces – Optional interface tagging list for feature extraction.

  • raise_if_gt_is_zero – Validate non-zero ground truth when gt_file is provided.

  • debug_plots_prefix – Plot filename prefix; defaults to skeleton stem.

  • return_frame – Include the ForSys frame under 'frame' when True.

  • verbose – Enable progress logging.

Returns:

Graph tensors and optional targets / predictions, same style as se_output_to_graph().

Return type:

dict

Models

stressnet.list_models() → list[str][source]

Return the names of all registered models.

stressnet.load_model(name: str | None = None, *, force_download: bool = False, device: str | None = None, **net_kwargs: Any) → tf.keras.Model[source]

Download (if needed) and return a ready-to-use StressNET model.

Parameters:

  • name – Model name as listed by list_models().

  • force_download – Re-download the weights even if they are already cached.

  • device – TensorFlow device string (e.g. '/GPU:0'). Defaults to GPU when available, otherwise CPU.

  • **net_kwargs – Extra keyword arguments forwarded to stressnet.build_stressnet().

Returns:

Compiled model with weights loaded, ready for inference.

Return type:

tf.keras.Model

stressnet.get_model_path(name: str) → Path[source]

Return the cached path for name, downloading it if necessary.

stressnet.clear_cache(name: str | None = None) → None[source]

Delete cached weights.

Parameters:

name – Model name to remove. Pass None to clear the entire cache.

stressnet.get_model(weights_path, points_per_edge: int = 9, disjoint_mode: bool = False, device: str | None = None, **net_kwargs) → tensorflow.keras.Model[source]

Build StressNET and optionally load weights on the chosen device.

Parameters:

  • weights_path – Path to an .h5 weights file (may be None only for random init / debugging).

  • points_per_edge – Number of sampled points per edge (must match training / checkpoint).

  • disjoint_mode – If True, build the disjoint-batched variant (Spektral disjoint mode).

  • device – TensorFlow device string; defaults to GPU when available.

  • **net_kwargs – Forwarded to build_stressnet().

Returns:

Compiled Keras model.

Return type:

tensorflow.keras.Model

stressnet.build_stressnet(load_weights_path: str | None = None, edge_n_vertices: int = 9, disjoint_mode: bool = False, **kwargs) → tensorflow.keras.models.Model[source]

Construct the StressNET graph neural network architecture.

Parameters:

  • load_weights_path – Optional path to load trained weights after building.

  • edge_n_vertices – Edge discretization count (controls input edge feature dimension).

  • disjoint_mode – Whether to include disjoint-mode graph index input.

  • **kwargs – Architecture hyperparameters (layer sizes, regularization, fine_tune_layers, etc.).

Returns:

Uncompiled or compiled model per internal Model factory defaults.

Return type:

tensorflow.keras.Model

Inference

stressnet.predict(model: tensorflow.keras.Model, data: dict[str, Any], device: str | None = None) → ndarray[source]

Run a single forward pass and return flattened stress predictions.

Parameters:

  • model – Compiled StressNET Keras model.

  • data – Graph batch with keys 'x', 'a' (sparse adjacency), 'e' (edge features).

  • device – TensorFlow device string. Defaults to GPU when available, else CPU.

Returns:

1-D array of predictions (length equals number of graph nodes for stress mode).

Return type:

numpy.ndarray

stressnet.predict_augmented(model: tensorflow.keras.Model, data: dict[str, Any], n_augmentations: int = 9, seed: int | None = 1337, device: str | None = None, mode: Literal['stress', 'embedding'] = 'stress') → ndarray[source]

Average predictions over random rotations of edge features (data augmentation).

Parameters:

  • model – Compiled StressNET Keras model.

  • data – Graph batch with keys 'x', 'a', 'e'.

  • n_augmentations – Number of additional random rotations (total runs = n_augmentations + 1).

  • seed – RNG seed for reproducibility.

  • device – TensorFlow device string.

  • mode – 'stress' rescales the mean prediction to unit mean; 'embedding' returns averaged embedding without that rescaling.

Returns:

Mean prediction vector across augmentations.

Return type:

numpy.ndarray

stressnet.calculate_metrics(y_true: ndarray, y_pred: ndarray) → tuple[float, float, float, float][source]

Correlation, R², MAPE (%), and composite Borges-style score.

Parameters:

  • y_true – Ground-truth targets (1-D).

  • y_pred – Model predictions (same shape as y_true).

Returns:

(pearson_r, r2, mape_percent, score).

Return type:

tuple[float, float, float, float]

stressnet.frame_with_predicted_tensions(frame: Frame, predicted_tensions: ndarray) → Frame[source]

Return a deep-copied frame with predicted tensions assigned to edges.

Parameters:

  • frame – Source ForSys frame (unchanged).

  • predicted_tensions – One scalar per big edge, in the same order as frame.get_big_edges().

Returns:

New frame instance; each small edge in a big edge receives the same tension.

Return type:

forsys.frames.Frame

Raises:

ValueError – If the number of predictions does not match the number of big edges.

Plotting

stressnet.plot_with_force(frame: Frame, filename: str | None = None, force_to_plot: Literal['stress', 'tension', 'gt', 'ground-truth'] | None = None, mirror_y: bool = False, figsize: tuple[float, float] = (10, 10), **kwargs) → None[source]

Plot the tissue graph, optionally coloring edges by tension or ground truth.

Parameters:

  • frame – ForSys frame with vertices/edges populated.

  • filename – If set, save the figure to this path (.png appended when extension missing).

  • force_to_plot – None draws discrete colors per internal big edge; 'stress' / 'tension' maps continuous edge tension; 'gt' / 'ground-truth' maps gt.

  • mirror_y – If True, invert the y-axis (image-style coordinates).

  • figsize – Matplotlib figure size in inches.

  • **kwargs – Common keys: cmap, cbar, cbar_step, cbar_params, title, plot_kwargs (merged into plt.plot for non-external edges).

Returns:

Displays interactively when filename is None; otherwise writes file and closes.

Return type:

None

Graph IO

stressnet.save_graph_data(dst_file: str, *, adj_mat: csr_matrix, node_features: ndarray, edge_features: ndarray, targets: ndarray | None = None, forsys_predictions: ndarray | None = None, compressed: bool = False, verbose: bool = False) → None[source]

Persist graph tensors (and optional labels) to an .npz file.

Parameters:

  • dst_file – Output path (.npz). Parent directories are created as needed.

  • adj_mat – Sparse CSR adjacency matrix.

  • node_features – Dense node feature matrix X.

  • edge_features – Dense edge feature matrix E.

  • targets – Optional ground-truth vector or matrix Y.

  • forsys_predictions – Optional ForSys baseline predictions stored as forsys_preds.

  • compressed – Use numpy.savez_compressed instead of numpy.savez.

  • verbose – Log a short summary when True.

stressnet.load_graph_data(file_path: str, include_forsys_predictions: bool = True, include_targets: bool = True, as_dict: bool = True)[source]

Load graph tensors produced by save_graph_data().

Parameters:

  • file_path – Path to an .npz archive written by this module.

  • include_forsys_predictions – If False, omit forsys_preds even when present in the file.

  • include_targets – If False, omit y even when present.

  • as_dict – If True (default), return a dict with keys 'a', 'x', 'e', and optionally 'y', 'forsys_preds'. If False, return a tuple in fixed order.

Returns:

Reconstructed adjacency, features, and optional arrays.

Return type:

dict or tuple

Notes

Tuple order when as_dict is False is (A, X, E[, Y][, forsys_predictions]) depending on the two include_* flags.