Ceren dev snaphot

.gitignore

@@ -24,4 +24,6 @@ Thumbs.db
 *.html
 *.png
 mlruns
-outputs
+outputs# pixi environments
+.pixi/*
+!.pixi/config.toml

RunDNAm_small.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+#SBATCH --job-name=Bulk_DNAmethylation
+#SBATCH --account=project_2015212
+#SBATCH --partition=gpu
+#SBATCH --gres=gpu:v100:1
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=160G
+#SBATCH --time=05:00:00
+#SBATCH --output=/scratch/project_2015212/ceren/runs/bulk/%x-%j.out
+#SBATCH --error=/scratch/project_2015212/ceren/runs/bulk/%x-%j.err
+
+set -euo pipefail
+
+module load tensorflow/2.18
+source /projappl/project_2015212/cavachon/envs/ceren/.venv/bin/activate
+
+export MLFLOW_TRACKING_URI="file:///scratch/project_2015212/ceren/mlruns"
+export OMP_NUM_THREADS=${SLURM_CPUS_PER_TASK}
+export MKL_NUM_THREADS=${SLURM_CPUS_PER_TASK}
+export PYTHONUNBUFFERED=1
+
+# Make sure key dirs exist
+mkdir -p /scratch/project_2015212/ceren/runs/bulk/embeddings
+mkdir -p /scratch/project_2015212/ceren/checkpoints
+
+cd /projappl/project_2015212/cavachon/CAVACHON
+
+python - << 'PY'
+from cavachon.workflow import Workflow
+CFG = "/projappl/project_2015212/cavachon/configs/ceren/DNAmethyl_small_run.yaml"
+wf = Workflow(CFG)
+wf.run()
+PY

RunDNAm_small2.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+#SBATCH --job-name=Bulk_DNAmethylation
+#SBATCH --account=project_2015212
+#SBATCH --partition=gpu
+#SBATCH --gres=gpu:v100:1
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=90G
+#SBATCH --time=01:00:00
+#SBATCH --output=/scratch/project_2015212/ceren/runs/bulk2/%x-%j.out
+#SBATCH --error=/scratch/project_2015212/ceren/runs/bulk2/%x-%j.err
+
+set -euo pipefail
+
+module load tensorflow/2.18
+source /projappl/project_2015212/cavachon/envs/ceren/.venv/bin/activate
+
+export MLFLOW_TRACKING_URI="file:///scratch/project_2015212/ceren/mlruns2"
+export OMP_NUM_THREADS=${SLURM_CPUS_PER_TASK}
+export MKL_NUM_THREADS=${SLURM_CPUS_PER_TASK}
+export PYTHONUNBUFFERED=1
+
+# Make sure key dirs exist
+mkdir -p /scratch/project_2015212/ceren/runs/bulk2/embeddings
+mkdir -p /scratch/project_2015212/ceren/checkpoints2
+
+cd /projappl/project_2015212/cavachon/CAVACHON
+
+python - << 'PY'
+from cavachon.workflow import Workflow
+CFG = "/projappl/project_2015212/cavachon/configs/ceren/DNAm_second.yaml"
+wf = Workflow(CFG)
+wf.run()
+PY

cavachon/dataloader/modifiers/__init__.py

@@ -1,6 +1,12 @@
-from .independent_bernoulli_data_modifier import (
-    IndependentBernoulliDataModifier as IndependentBernoulliDataModifier,
-)
-from .independent_zero_inflated_negative_binomial_data_modifier import (
-    IndependentZeroInflatedNegativeBinomialDataModifier as IndependentZeroInflatedNegativeBinomialDataModifier,
-)
+from .independent_bernoulli_data_modifier import (
+    IndependentBernoulliDataModifier as IndependentBernoulliDataModifier,
+)
+from .independent_zero_inflated_negative_binomial_data_modifier import (
+    IndependentZeroInflatedNegativeBinomialDataModifier as IndependentZeroInflatedNegativeBinomialDataModifier,
+)
+from .multivariate_normal_diag_data_modifier import (
+    MultivariateNormalDiagDataModifier as MultivariateNormalDiagDataModifier,
+)
+from .studentt_data_modifier import (
+    StudenttDataModifier as StudenttDataModifier,
+)

cavachon/dataloader/modifiers/multivariate_normal_diag_data_modifier.py

@@ -0,0 +1,74 @@
+import functools
+from typing import Any, Mapping
+
+import tensorflow as tf
+
+from cavachon.environment.constants import Constants
+from cavachon.layers.modifiers.to_dense import ToDense
+
+
+class MultivariateNormalDiagDataModifier(tf.keras.Model):
+    """MultivariateNormalDiagDataModifier
+
+    Modifiers for the modality which follows a MultivariateNormalDiag
+    distribution (Normal distribution with diagonal covariance).
+    The instance will be used right after the tf.data.Dataset is
+    created using the DataLoader.
+
+    Attributes
+    ----------
+    modality_names: str
+        modality name.
+
+    modality_key: str
+        the key used to access the mapping of data created from
+        tf.data.Dataset. Defaults to `modality_name`_matrix.
+
+    modifiers: List[tf.keras.layers.Layer]
+        list of modifiers that will be applied to the data created from
+        tf.data.Dataset. Defaults to [ToDense].
+
+    See Also
+    --------
+    DataLoader: used to create tf.data.Dataset from MuData.
+
+    """
+
+    def __init__(self, modality_name: str):
+        """Constructor for MultivariateNormalDiag data modifier
+
+        Parameters
+        ----------
+        modality_name: str
+            the name of modality that needs to be processed.
+        """
+        super().__init__()
+        self.modality_name: str = modality_name
+        self.modality_key: str = f"{modality_name}_{Constants.TENSOR_NAME_X}"
+        # For continuous normalized data (CNV, normalized RNA, etc.)
+        # we only need to convert sparse matrices to dense tensors
+        self.modifiers = [ToDense(self.modality_key)]
+
+    def call(self, inputs: Mapping[Any, tf.Tensor], training=None, mask=None):
+        """Process the data created from tf.data.Dataset.
+
+        Parameters
+        ----------
+        inputs:
+            Mapping of tf.Tensor, where the keys contain
+            self.modality_key.
+
+        training: bool, optional
+            Not used (kept for tf.keras.Model API).
+
+        mask: tf.Tensor, optional
+            Not used (kept for tf.keras.Model API).
+
+        Returns
+        -------
+        Mapping[Any, tf.Tensor]
+            processed data.
+
+        """
+        modifiers = self.modifiers
+        return functools.reduce(lambda x, modifier: modifier(x), modifiers, inputs)

cavachon/dataloader/modifiers/studentt_data_modifier.py

@@ -0,0 +1,17 @@
+import functools
+from typing import Any, Mapping
+
+import tensorflow as tf
+
+from cavachon.environment.constants import Constants
+from cavachon.layers.modifiers.to_dense import ToDense
+
+
+class StudenttDataModifier(tf.keras.Model):
+    def __init__(self, modality_name: str):
+        super().__init__()
+        self.modality_key = f"{modality_name}_{Constants.TENSOR_NAME_X}"
+        self.modifiers = [ToDense(self.modality_key)]
+
+    def call(self, inputs: Mapping[Any, tf.Tensor], **kwargs):
+        return functools.reduce(lambda x, mod: mod(x), self.modifiers, inputs)

cavachon/distributions/__init__.py

@@ -8,3 +8,6 @@
 from .multivariate_normal_diag_distribution import (
     MultivariateNormalDiagDistribution as MultivariateNormalDiagDistribution,
 )
+from .studentt_distribution import (
+    StudenttDistribution as StudenttDistribution,
+)

cavachon/distributions/distribution.py

@@ -1,4 +1,4 @@
-from abc import ABC, abstractclassmethod
+from abc import ABC, abstractmethod
 from typing import Mapping, Union
 
 import tensorflow as tf
@@ -11,7 +11,8 @@ class Distribution(ABC):
 
     """
 
-    @abstractclassmethod
+    @classmethod
+    @abstractmethod
     def from_parameterizer_output(
         cls, params: Union[tf.Tensor, Mapping[str, tf.Tensor]], **kwargs
     ) -> tfp.distributions.Distribution:

cavachon/distributions/studentt_distribution.py

@@ -0,0 +1,29 @@
+import tensorflow as tf
+import tensorflow_probability as tfp
+
+from cavachon.distributions.distribution import Distribution
+
+
+class StudenttDistribution(Distribution, tfp.distributions.StudentT):
+    """StudentT distribution for continuous data with heavy tails (e.g. CNV)."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @classmethod
+    def from_parameterizer_output(cls, params: tf.Tensor, **kwargs):
+        """
+        Creates distribution from a single tensor.
+        The last dimension is split into 3: loc, scale, and df.
+        """
+        # Split into 3 equal parts
+        loc, scale_raw, df_raw = tf.split(params, 3, axis=-1)
+
+        # Scale (sigma) must be positive
+        scale = tf.math.softplus(scale_raw) + 1e-7
+
+        # Degrees of Freedom (nu) must be > 0.
+        # Adding 2.0 ensures the variance is mathematically defined (> 2).
+        df = tf.math.softplus(df_raw) + 2.0
+
+        return cls(df=df, loc=loc, scale=scale, **kwargs)

cavachon/layers/parameterizers/__init__.py

@@ -13,3 +13,9 @@
 from .multivariate_normal_diag_sampler import (
     MultivariateNormalDiagSampler as MultivariateNormalDiagSampler,
 )
+from .studentt_parameterizer_layer import (
+    StudenttParameterizerLayer as StudenttParameterizerLayer,
+)
+from .studentt_sampler import (
+    StudenttSampler as StudenttSampler,
+)