susan.project.SubtomoAvgN2N

Warning

Experimental / in development. API, default hyper-parameters, and per-iteration behaviour are not yet stable; results have not been broadly validated. Do not rely on this class for production workflows.

class susan.project.SubtomoAvgN2N.SubtomoAvgN2N(prj_name, box_size=None, n2n_arch='convnet', n2n_depth=6, n2n_n_feat=32, n2n_loss='l1')[source]

Bases: SubtomoAvgSched

Subtomogram averaging with Noise2Noise denoising each iteration.

Warning

Experimental / in development. API and default hyper-parameters may change; results have not been broadly validated. Not recommended for production workflows.

Extends SubtomoAvgSched by overriding run_postprocessing() to train a Noise2Noise model on the freshly reconstructed half-maps and replace the reference map with the denoised output.

The model and trainer are not stored as persistent instance attributes. Their construction parameters are stored, and the objects are re-instantiated each call so that GPU memory is released after training.

Parameters:

  • prj_name (str) – Project directory.

  • box_size (int, optional) – Box size in pixels. Reads info.prjtxt when omitted.

  • n2n_arch (str, optional) – Backbone architecture: 'convnet' or 'unet'. Default: 'convnet'.

  • n2n_depth (int, optional) – For 'convnet': total number of blocks. For 'unet': number of encoder/decoder levels. Default: 6.

  • n2n_n_feat (int, optional) – Number of feature maps per layer. Default: 32.

  • n2n_loss (str, optional) – 'l1' (MAE) or 'l2' (MSE). Default: 'l1'.

n2n_arch, n2n_depth, n2n_n_feat, n2n_loss

Architecture / loss parameters saved for model recreation.

n2n_scratch_lr, n2n_scratch_epochs, n2n_scratch_n_pairs

Training parameters used when no previous model exists.

Type:

float / int

n2n_finetune_lr, n2n_finetune_epochs, n2n_finetune_n_pairs

Training parameters used for fine-tuning from a previous model.

Type:

float / int

n2n_batch_size

Pairs per gradient step. Default: 1.

Type:

int

n2n_print_every

Print loss every N epochs (0 to silence). Default: 10.

Type:

int

n2n_mask

Path to an MRC file used as mask for set_size_mask(). The file is read each time a new VolumePairs buffer is allocated. If None, set_size() is used instead with padding=0.

Type:

str or None

n2n_extra_pad

Extra padding for set_size_mask(). Default: 0.

Type:

int

n2n_num_entries

Passed to populate() as num_entries. None fills to capacity.

Type:

int or None

n2n_free_gpu

If True (default), delete model and trainer after each call to train_and_denoise() and empty the CUDA cache.

Type:

bool

n2n_ptcls_modifier

Optional (ptcls, tomos) -> None hook applied to particles in run_postprocessing() (automated loop only) after the standard tilt-limit logic. None means no extra modification.

Type:

callable or None

Manual Loop

create_training_dataset(ptcls, fine_tune=True)[source]

Build (or rebuild) the internal VolumePairs training buffer.

The caller is responsible for any particle modifications (tilt limiting, flag zeroing, etc.) before passing ptcls.

Parameters:

  • ptcls (Particles) – Particle stack to reconstruct from. Modified temporarily during population but restored on exit.

  • fine_tune (bool, optional) – If True (default) use n2n_finetune_n_pairs. If False use n2n_scratch_n_pairs.

train_and_denoise(ite, fine_tune=True)[source]

Train the N2N model on the current dataset and denoise the map.

Loads weights from the previous iteration (or initial model if available, else fresh weights), trains, saves the model to <ite_dir>/model.pth, denoises the map, and returns the inverted denoised volume ready to write to disk.

create_training_dataset() must be called before this method.

Parameters:

  • ite (int) – Current iteration number.

  • fine_tune (bool, optional) – If True (default) use finetune lr/epochs. If False use scratch lr/epochs.

Returns:

Denoised map, sign-inverted and padded to full box size.

Return type:

numpy.ndarray