add class methods to load etomo and aretomo data directly

[McHaillet]

big +1 on class methods from my side

Tomogram.from_etomo_directory()
Tomogram.from_aretomo_aln()

Originally posted by @alisterburt in #14 (comment)

[davidetorre99]

Would something like this work?
This approach would directly use etomofiles and alnfile to parse the alignment files into their respective DataFrames and then output the class ready for reconstruction:

@classmethod
def from_aretomo_aln(
    cls,
    aln_path: Path | str,
    tilt_stack_path: Path | str,
    device: torch.device | str = "cpu",
) -> "Tomogram":
    """Initialize Tomogram from AreTomo alignment file and tilt stack."""

    aln_path = Path(aln_path)
    tilt_stack_path = Path(tilt_stack_path)
    df = alnfile.read(aln_path)
    corrected_shifts_xy = df[["tx", "ty"]].to_numpy()
    corrected_shifts_yx = corrected_shifts_xy[:, ::-1].copy() 
    tilt_stack_full = mrcfile.read(tilt_stack_path)
    included_indices = df["sec"].values - 1 
    tilt_stack = tilt_stack_full[included_indices]
    tilt_stack = tilt_stack.astype(np.float32)
    tilt_stack -= np.mean(tilt_stack, axis=(-2, -1), keepdims=True)
    tilt_stack /= np.std(tilt_stack, axis=(-2, -1), keepdims=True)
    return cls(
        images=tilt_stack,
        tilt_angles=df["tilt"].to_numpy(),
        tilt_axis_angle=df["rot"].to_numpy(),
        sample_translations=corrected_shifts_yx,
        device=device,
    )

@classmethod
def from_etomo_directory(
    cls,
    etomo_dir: Path | str,
    normalize: bool = True,
    device: torch.device | str = "cpu",
) -> "Tomogram":
    """Initialize Tomogram from eTomo directory."""

    etomo_dir = Path(etomo_dir)
    df = etomofiles.read(etomo_dir)
    df = df.loc[~df["excluded"]].reset_index(drop=True)
    xf = etomofiles.df_to_xf(df, yx=True)
    m, shifts = xf[:, :, :2], xf[:, :, 2]
    corrected_shifts = -np.einsum("nji,nj->ni", m, shifts)
    corrected_shifts = np.ascontiguousarray(corrected_shifts)
    tilt_stack_path = etomo_dir / df.image_path[0].replace("[0]", "")
    tilt_stack_full = mrcfile.read(tilt_stack_path)
    tilt_stack = tilt_stack_full[df.idx_tilt.to_numpy()]
    tilt_stack = tilt_stack.astype(np.float32)
    tilt_stack -= np.mean(tilt_stack, axis=(-2, -1), keepdims=True)
    tilt_stack /= np.std(tilt_stack, axis=(-2, -1), keepdims=True)

    return cls(
        images=tilt_stack,
        tilt_angles=df.tlt.to_numpy(),
        tilt_axis_angle=df.tilt_axis_angle.to_numpy(),
        sample_translations=corrected_shifts,
        device=device,
    )

It could then be used like:

import mrcfile
from pathlib import Path
from torch_tomogram import Tomogram

ETOMO_DIR = Path("/path/to/etomo/dir")

tilt_series = Tomogram.from_etomo_directory(
    etomo_dir=ETOMO_DIR,
    device="cpu"
)

# Reconstruct tomogram
tomogram = tilt_series.reconstruct_tomogram((100, 480, 380), 128)

# Save tomogram
mrcfile.write(
    ETOMO_DIR / "tt_rec.mrc",
    tomogram.numpy(),
    overwrite=True,
    voxel_size=10

)

Or would you rather pass the dataframe to the classmethods?