adding easyconfigs: tomoDRGN-1.0.0-foss-2023a-CUDA-12.1.1.eb and patc…

…hes: tomoDRGN-1.0.0_use-PyTorch-2.1.patch

easybuild/easyconfigs/t/tomoDRGN/tomoDRGN-1.0.0-foss-2023a-CUDA-12.1.1.eb

@@ -0,0 +1,92 @@
+# Thomas Hoffmann, EMBL Heidelber, [email protected], 2024/12
+easyblock = 'PythonBundle'
+
+name = 'tomoDRGN'
+version = '1.0.0'
+versionsuffix = '-CUDA-%(cudaver)s'
+
+homepage = 'https://github.com/bpowell122/tomodrgn/'
+description = """CryoDRGN has proven a powerful deep learning method for heterogeneity analysis  
+in single particle cryo-EM. In particular, the method models a continuous       
+distribution over 3D structures by using a Variational Auto-Encoder (VAE) based 
+architecture to generate a reconstruction voxel-by-voxel once given a fixed     
+coordinate from a continuous learned latent space.
+
+TomoDRGN extends the cryoDRGN framework to cryo-ET by learning heterogeneity    
+from datasets in which each particle is sampled by multiple projection images at
+different stage tilt angles. For cryo-ET samples imaging particles in situ,     
+tomoDRGN therefore enables continuous heterogeneity analysis at a single        
+particle level within the native cellular environment. This new type of input   
+necessitates modification of the cryoDRGN architecture, enables tomography-     
+specific processing opportunities (e.g. dose weighting for loss weighting and   
+efficient voxel subset evaluation during training), and benefits from           
+tomography-specific interactive visualizations.
+"""
+
+toolchain = {'name': 'foss', 'version': '2023a'}
+
+builddependencies = [
+    ('nodejs', '18.17.1'),
+    ('poetry', '1.5.1')
+]
+dependencies = [
+    ('Python', '3.11.3'),
+    ('CUDA', '12.1.1', '', SYSTEM),
+    ('JupyterNotebook', '7.0.2'),
+    ('SciPy-bundle', '2023.07'),
+    ('PyTorch', '2.1.2', versionsuffix),
+    ('scikit-learn', '1.3.1'),
+    ('matplotlib', '3.7.2'),
+    ('umap-learn', '0.5.5'),
+    ('plotly.py', '5.16.0'),
+    ('healpy', '1.16.6'),
+    ('Seaborn', '0.13.2'),
+    ('IPython', '8.14.0'),
+    ('einops', '0.7.0'),
+    ('adjustText', '0.7.3'),
+]
+
+use_pip = True
+sanity_pip_check = True
+
+exts_list = [
+    ('fastcluster', '1.2.6', {
+        'checksums': ['aab886efa7b6bba7ac124f4498153d053e5a08b822d2254926b7206cdf5a8aa6'],
+    }),
+    ('importlib_resources', '5.12.0', {
+        'checksums': ['4be82589bf5c1d7999aedf2a45159d10cb3ca4f19b2271f8792bc8e6da7b22f6'],
+    }),
+    ('torchinfo', '1.8.0', {
+        'checksums': ['72e94b0e9a3e64dc583a8e5b7940b8938a1ac0f033f795457f27e6f4e7afa2e9'],
+    }),
+    (name, version, {
+        'patches': ['%(name)s-%(version)s_use-PyTorch-2.1.patch'],
+        'preinstallopts': "export SETUPTOOLS_SCM_PRETEND_VERSION=%(version)s &&",
+        'source_urls': ['https://github.com/bpowell122/tomodrgn/archive/refs/tags/'],
+        'sources': {'download_filename': 'v%(version)s.tar.gz', 'filename': '%(name)s-%(version)s.tar.gz'},
+        'checksums': [
+            {'tomoDRGN-1.0.0.tar.gz': '8b243d17af655c5088b30b4a9476072eb79e144b0a9ec1be549457c11abdd2d7'},
+            {'tomoDRGN-1.0.0_use-PyTorch-2.1.patch':
+             '5c7c1bec937ca245c999d2815f9345882950ae3da23432704efb7fb3e61f512b'},
+        ],
+    }),
+]
+
+postinstallcmds = [
+    # required?
+    'mkdir %(installdir)s/share/jupyter/labextensions -p',
+    'export JUPYTERLAB_DIR=%(installdir)s/share/jupyter/labextensions '
+    '&& jupyter labextension install jupyterlab-datawidgets'
+]
+
+
+sanity_check_paths = {
+    'files': ['bin/%(namelower)s'],
+    'dirs': ['share/jupyter/labextensions']
+}
+
+modextravars = {'JUPYTERLAB_DIR': '%(installdir)s/share/jupyter/labextensions'}  # required?
+
+sanity_check_commands = ['%(namelower)s -h']
+
+moduleclass = 'chem'

easybuild/easyconfigs/t/tomoDRGN/tomoDRGN-1.0.0_use-PyTorch-2.1.patch

@@ -0,0 +1,66 @@
+# Downgrade to PyTorch 2.1.2 (CUDA):
+#   relax requirement torch >= 2.3 to 2.1
+#   replace any occurrence torch.amp.GradScaler by conditional torch.cuda.amp / torch.cpu.amp
+diff -ru tomodrgn-1.0.0/pyproject.toml tomodrgn-1.0.0_relax_requirements/pyproject.toml
+--- tomodrgn-1.0.0/pyproject.toml	2024-11-15 21:15:46.000000000 +0100
++++ tomodrgn-1.0.0_relax_requirements/pyproject.toml	2024-12-16 15:36:11.187436498 +0100
+@@ -31,7 +31,7 @@
+     "scikit-learn",
+     "scipy>=1.3.1",
+     "seaborn",
+-    "torch>=2.3",
++    "torch>=2.1",
+     "torchinfo",
+     "typing_extensions>=3.7.4",
+     "umap-learn",
+diff -ru tomodrgn-1.0.0/tomodrgn/commands/train_nn.py tomodrgn-1.0.0_relax_requirements/tomodrgn/commands/train_nn.py
+--- tomodrgn-1.0.0/tomodrgn/commands/train_nn.py	2024-11-15 21:15:46.000000000 +0100
++++ tomodrgn-1.0.0_relax_requirements/tomodrgn/commands/train_nn.py	2024-12-17 10:58:56.367688173 +0100
+@@ -145,7 +145,7 @@
+
+
+ def train_batch(model: FTPositionalDecoder | DataParallelPassthrough,
+-                scaler: torch.amp.GradScaler,
++                scaler: torch.cuda.amp.GradScaler,
+                 optim: torch.optim.Optimizer,
+                 lat: Lattice,
+                 batch_images: torch.Tensor,
+@@ -402,7 +402,7 @@
+     # Mixed precision training with AMP
+     use_amp = not args.no_amp
+     flog(f'AMP acceleration enabled (autocast + gradscaler) : {use_amp}')
+-    scaler = torch.amp.GradScaler(device=device.type, enabled=use_amp)
++    scaler = torch.cuda.amp.GradScaler(enabled=use_amp) if device.type == 'cuda' else torch.cpu.amp.GradScaler(enabled=use_amp) 
+     if use_amp:
+         if not args.batch_size % 8 == 0:
+             flog('Warning: recommended to have batch size divisible by 8 for AMP training')
+diff -ru tomodrgn-1.0.0/tomodrgn/commands/train_vae.py tomodrgn-1.0.0_relax_requirements/tomodrgn/commands/train_vae.py
+--- tomodrgn-1.0.0/tomodrgn/commands/train_vae.py	2024-11-15 21:15:46.000000000 +0100
++++ tomodrgn-1.0.0_relax_requirements/tomodrgn/commands/train_vae.py	2024-12-17 10:59:45.399859026 +0100
+@@ -118,7 +118,7 @@
+
+ def train_batch(*,
+                 model: TiltSeriesHetOnlyVAE | DataParallelPassthrough,
+-                scaler: torch.amp.GradScaler,
++                scaler: torch.cuda.amp.GradScaler,
+                 optim: torch.optim.Optimizer,
+                 lat: Lattice,
+                 batch_images: torch.Tensor,
+@@ -442,7 +442,7 @@
+
+ def save_checkpoint(*,
+                     model: TiltSeriesHetOnlyVAE | DataParallelPassthrough,
+-                    scaler: torch.amp.GradScaler,
++                    scaler: torch.cuda.amp.GradScaler,
+                     optim: torch.optim.Optimizer,
+                     epoch: int,
+                     z_mu_train: np.ndarray,
+@@ -692,7 +692,7 @@
+     # Mixed precision training with AMP
+     use_amp = not args.no_amp
+     flog(f'AMP acceleration enabled (autocast + gradscaler) : {use_amp}')
+-    scaler = torch.amp.GradScaler(device=device.type, enabled=use_amp)
++    scaler = torch.cuda.amp.GradScaler(enabled=use_amp) if device.type == 'cuda' else torch.cpu.amp.GradScaler(enabled=use_amp) 
+     if use_amp:
+         if not args.batch_size % 8 == 0:
+             flog('Warning: recommended to have batch size divisible by 8 for AMP training')