### Install PySCENIC from GitHub

Source: https://tuftsbcb.github.io/RegDiffusion/_sources/downstream_with_pyscenic.md

Install the recommended GitHub version of pyscenic into your current working environment.

```bash
>>> pip install git+https://github.com/aertslab/pySCENIC
```

--------------------------------

### Install RegDiffusion

Source: https://tuftsbcb.github.io/RegDiffusion/index.html

Install the regdiffusion package using pip.

```bash
pip install regdiffusion
```

--------------------------------

### Registering a Buffer

Source: https://tuftsbcb.github.io/RegDiffusion/_autosummary/regdiffusion.models.RegDiffusion.html

Demonstrates how to register a buffer with a module, which is state that is not a model parameter but should be saved. Example shows registering 'running_mean'.

```python
>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))
```

--------------------------------

### Example of accessing state_dict keys

Source: https://tuftsbcb.github.io/RegDiffusion/_autosummary/regdiffusion.models.RegDiffusion.html

Demonstrates how to access the keys of a module's state dictionary. This is useful for understanding the parameters and buffers managed by the module.

```python
>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']
```

--------------------------------

### Iterating Through Parameters

Source: https://tuftsbcb.github.io/RegDiffusion/_autosummary/regdiffusion.models.RegDiffusion.html

Shows how to iterate over all parameters of a module, typically used when passing parameters to an optimizer. Includes examples of parameter types and sizes.

```python
>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
```

--------------------------------

### Iterating Through Named Parameters

Source: https://tuftsbcb.github.io/RegDiffusion/_autosummary/regdiffusion.models.RegDiffusion.html

Provides an example of iterating over module parameters, yielding their names and the parameters themselves. Useful for inspecting specific parameters like 'bias'.

```python
>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())
```

--------------------------------

### Create and Activate Conda Environment for PySCENIC

Source: https://tuftsbcb.github.io/RegDiffusion/_sources/downstream_with_pyscenic.md

Set up a dedicated conda environment with specific package versions to avoid conflicts when installing and using pyscenic.

```bash
>>> conda create -y -n pyscenic-env python=3.10
>>> conda activate pyscenic-env
>>> pip install pyscenic==0.12.1 numpy==1.23.5 pandas==1.3.5 dask==2023.10.0 matplotlib decorator==4.3.0
>>> conda deactivate
```

--------------------------------

### Accessing Named Buffers in a Module

Source: https://tuftsbcb.github.io/RegDiffusion/_autosummary/regdiffusion.models.RegDiffusion.html

Provides an example of iterating over named buffers within a module. This is useful for accessing specific buffers like 'running_var' and checking their sizes.

```python
>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

```

--------------------------------

### Initialize and Apply Weights to a Sequential Model

Source: https://tuftsbcb.github.io/RegDiffusion/_autosummary/regdiffusion.models.RegDiffusion.html

Demonstrates how to apply a custom weight initialization function to all submodules within a `nn.Sequential` model. This is useful for setting initial parameter values.

```python
>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

```

--------------------------------

### Get Gene Embeddings

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/models/regdiffusion.html

Retrieves the gene embedding data as a NumPy array.

```python
def get_gene_emb(self):
        return self.gene_emb[0].gene_emb.data.cpu().detach().numpy()
```

--------------------------------

### Get Sampled Adjacency Matrix

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/models/regdiffusion.html

Retrieves the sampled adjacency matrix values, detached and ready for use.

```python
@torch.no_grad()
    def get_sampled_adj_(self):
        return self.get_adj_()[
            self.sampled_adj_row_nonparam, self.sampled_adj_col_nonparam
        ].detach()
```

--------------------------------

### Get Cleaned Adjacency Matrix

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/models/regdiffusion.html

Returns a cleaned adjacency matrix after applying soft thresholding and removing the diagonal.

```python
def get_adj_(self):
        mask = 1 - torch.eye(self.n_gene, device=self.adj_A.device)
        return self.soft_thresholding(self.adj_A, self.gene_reg_norm/2) * mask
```

--------------------------------

### Build Datasets for Training and Validation

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/trainer.html

Initializes training and validation datasets and dataloaders using sparse expression data. Uses RandomSampler for training with replacement and a fixed batch size.

```python
self.train_dataset = _SparseExpressionDataset(
            exp_array, cell_types,
            cell_min_f32, cell_range_f32, gene_mean, gene_std,
            train_indices)
        train_sampler = torch.utils.data.RandomSampler(
            self.train_dataset, replacement=True, num_samples=batch_size,
            generator=self._sampler_generator)
        self.train_dataloader = DataLoader(
            self.train_dataset,
            sampler=train_sampler,
            batch_size=batch_size,
            drop_last=True)

        self.val_dataset = _SparseExpressionDataset(
            exp_array, cell_types,
            cell_min_f32, cell_range_f32, gene_mean, gene_std,
            val_indices)
        self.val_dataloader = DataLoader(
            self.val_dataset,
            shuffle=False,
            batch_size=batch_size,
            drop_last=False)
```

--------------------------------

### Initialize RegDiffusion Trainer

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/trainer.html

Initializes the RegDiffusion trainer with specified hyperparameters, device, and logger. Handles device selection and raises an exception for unsupported devices like Apple Silicon's MPS. Sets up the diffusion schedule and prepares data based on its format (sparse or dense).

```python
if device == 'mps':
            raise Exception("We noticed unreliable training behavior on", 
                            "Apple's silicon. Consider using other devices.")
        elif device.startswith('cuda'):
            if not torch.cuda.is_available():
                print(
                    "You specified cuda as your computing device but apprently", 
                    "it's not available. Setting device to cpu for now. ")
                device = 'cpu'
        self.device = device
        self.hp['device'] = device
        
        # Logger ---------------------------------------------------------------
        if logger is None:
            self.logger = LightLogger()
        else:
            self.logger = logger
        self.note_id = self.logger.start()
        
        # Define diffusion schedule
        self.betas = linear_beta_schedule(T, start_noise, end_noise)
        self.alphas = 1. - self.betas
        alpha_bars = torch.cumprod(self.alphas, axis=0)
        self.mean_schedule = torch.sqrt(alpha_bars).to(device)
        self.std_schedule = torch.sqrt(1. - alpha_bars).to(device)
    
        # Prepare Data ---------------------------------------------------------
        if (exp_array.shape[0] < batch_size):
            warnings.warn(
                "Batch size needs to be smaller than the number of cells. "
            )
        if cell_types is None:
            cell_types = np.zeros(exp_array.shape[0], dtype=int)
        self.n_celltype = len(np.unique(cell_types))
        n_cell, n_gene = exp_array.shape
        self.n_cell = n_cell
        self.n_gene = n_gene

        self.evaluator = evaluator

        if sp.issparse(exp_array):
            self._prepare_sparse_data(
                exp_array, cell_types, batch_size,
                train_split, train_split_seed)
        else:
            self._prepare_dense_data(
                exp_array, cell_types, batch_size,
                train_split, train_split_seed)
    
        # Setup Model ----------------------------------------------------------
        gene_reg_norm = 1/(n_gene-1)
        ModelClass = RegDiffusionME if memory_efficient else RegDiffusion
        self.model = ModelClass(
            n_gene=n_gene,
            time_dim=time_dim,
            n_celltype=self.n_celltype,
            celltype_dim = celltype_dim,
            hidden_dims=hidden_dims,
            adj_dropout=adj_dropout,
            init_coef=init_coef
        )
    
        # Setup optimizer ------------------------------------------------------
        if lr_adj is None:
            lr_adj = gene_reg_norm/50
            self.hp['lr_adj'] = lr_adj
        adj_params = []
        non_adj_params = []
        for name, param in self.model.named_parameters():
            if name.endswith('adj_A'):
                adj_params.append(param)
            else:
                if not name.endswith('_nonparam'):
                    non_adj_params.append(param)
        self.opt = torch.optim.Adam(
            [{'params': non_adj_params}, {'params': adj_params}], 
            lr=lr_nn, 
            weight_decay=weight_decay_nn, betas=[0.9, 0.99]
        )
        self.opt.param_groups[0]['lr'] = lr_nn
        self.opt.param_groups[1]['lr'] = lr_adj
        self.opt.param_groups[1]['weight_decay'] = weight_decay_adj
    
        self.model.to(device)
        if self.device.startswith('cuda') and compile:
            self.original_model = self.model
            self.model = torch.compile(self.model)
        self.total_time_cost=0
        self.losses_on_gene=None
        self.model_name='RegDiffusionME' if memory_efficient else 'RegDiffusion'

        # AMP setup
        self.use_amp = use_amp
        self.amp_device_type = 'cuda' if device.startswith('cuda') else 'cpu'
```

--------------------------------

### Get Adjacency Matrix as NumPy Array

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/models/regdiffusion.html

Retrieves the adjacency matrix, detaches it, moves it to CPU, and converts it to a NumPy array of float16.

```python
def get_adj(self):
        adj = self.get_adj_().detach().cpu().numpy() / self.gene_reg_norm
        return adj.astype(np.float16)
```

--------------------------------

### Prepare Sparse Data with On-the-Fly Normalization

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/trainer.html

Processes sparse expression data to compute normalization statistics and create a dataset that normalizes data on-the-fly per sample. This method is memory-efficient for large datasets.

```python
    def _prepare_sparse_data(self, exp_array, cell_types, batch_size,
                             train_split, train_split_seed):
        """Compute normalization stats from sparse matrix in chunks, then build
        a _SparseExpressionDataset that normalizes on-the-fly per sample.

        This never materializes the full dense normalized matrix, so memory
        usage stays proportional to (sparse matrix size + one batch of dense
        rows) rather than (n_cell * n_gene * 4 bytes).
        """
        n_cell, n_gene = exp_array.shape
        exp_array = sp.csr_matrix(exp_array)  # ensure CSR for fast row slicing

        # --- Per-cell min/max in chunks ---
        chunk_size = min(10000, n_cell)
        cell_min = np.empty(n_cell, dtype=np.float64)
        cell_max = np.empty(n_cell, dtype=np.float64)
        for start in range(0, n_cell, chunk_size):
            end = min(start + chunk_size, n_cell)
            chunk = exp_array[start:end]
            cell_min[start:end] = np.asarray(
                chunk.min(axis=1).todense()).ravel()
            cell_max[start:end] = np.asarray(
                chunk.max(axis=1).todense()).ravel()

        cell_range = cell_max - cell_min
        valid_mask = cell_range > 0
        n_zero_cells = (~valid_mask).sum()
        if n_zero_cells > 0:
            warnings.warn(
                f'{n_zero_cells} cells are removed from analysis where no '
                'genes are expressed.')

        # --- Per-gene mean/std of min-max-normalized data (chunked) ---
        n_valid = valid_mask.sum()
        gene_sum = np.zeros(n_gene, dtype=np.float64)
        gene_sum_sq = np.zeros(n_gene, dtype=np.float64)

        for start in range(0, n_cell, chunk_size):
            end = min(start + chunk_size, n_cell)
            chunk_valid = valid_mask[start:end]
            if not chunk_valid.any():
                continue
            # Only materialize valid rows in this chunk
            chunk_dense = exp_array[start:end][chunk_valid].toarray().astype(
                np.float64)
            c_min = cell_min[start:end][chunk_valid][:, None]
            c_range = cell_range[start:end][chunk_valid][:, None]
            chunk_norm = (chunk_dense - c_min) / c_range
            gene_sum += chunk_norm.sum(axis=0)
            gene_sum_sq += (chunk_norm ** 2).sum(axis=0)

        gene_mean = (gene_sum / n_valid).astype(np.float32)
        gene_var = (gene_sum_sq / n_valid) - gene_mean.astype(np.float64) ** 2
        gene_std = np.sqrt(np.maximum(gene_var, 0)).astype(np.float32)

        n_zero_genes = (gene_std == 0).sum()
        if n_zero_genes > 0:
            raise ValueError(
                f'{n_zero_genes} genes have 0 variance. Please remove these '
                'genes from your data.')

        # --- Train/validation split (only among valid cells) ---
        valid_indices = np.where(valid_mask)[0]
        random_state = np.random.RandomState(train_split_seed)
        split_vals = random_state.rand(len(valid_indices))
        train_indices = valid_indices[split_vals <= train_split]
        val_indices = valid_indices[split_vals > train_split]

        cell_min_f32 = cell_min.astype(np.float32)
        cell_range_f32 = cell_range.astype(np.float32)
```

--------------------------------

### RegDiffusion Trainer Initialization

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/trainer.html

Initializes the RegDiffusion trainer with various hyperparameters. This includes settings for learning rates, regularization, batch size, training steps, device configuration, and optional components for evaluation and logging. It also handles seeding for reproducibility.

```python
def __init__(
        self,
        exp_array,
        cell_types=None,
        T=5000,
        start_noise=0.0001,
        end_noise=0.02,
        time_dim=64,
        celltype_dim=4,
        hidden_dims=[16, 16, 16],
        init_coef = 5,
        lr_nn=1e-3,
        lr_adj=None,
        weight_decay_nn=0.1,
        weight_decay_adj = 0.01,
        sparse_loss_coef=0.25,
        adj_dropout=0.30,
        batch_size=128,
        n_steps=1000,
        train_split=1.0,
        train_split_seed=123,
        device='cuda',
        compile=False,
        evaluator=None,
        eval_on_n_steps=100,
        logger=None,
        gradient_accumulation=False,
        use_amp=False,
        memory_efficient=False,
        seed=None
    ):
        hp = locals()
        del hp['exp_array']
        del hp['cell_types']
        del hp['logger']
        self.hp = hp

        # Seed for reproducibility
        if seed is not None:
            set_seed(seed)
        self._sampler_generator = torch.Generator()
        if seed is not None:
            self._sampler_generator.manual_seed(seed)
```

--------------------------------

### Filter Genes with Scanpy

Source: https://tuftsbcb.github.io/RegDiffusion/large_networks.html

Reduce memory usage by filtering genes. This example removes genes not expressed in at least 10 cells using `scanpy.pp.filter_genes`.

```python
import scanpy as sc

# Remove genes expressed in fewer than 10 cells
sc.pp.filter_genes(adata, min_cells=10)

```

--------------------------------

### Train RegDiffusion on CPU

Source: https://tuftsbcb.github.io/RegDiffusion/_sources/large_networks.md

Instantiate and train the RegDiffusion model using the CPU. Set the device to 'cpu' and specify the number of training steps.

```python
trainer = rd.RegDiffusionTrainer(
    exp_array,
    device='cpu',
    n_steps=1000
)
trainer.train()
```

--------------------------------

### Get Adjacency Matrix

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/trainer.html

Retrieves the scaled adjacency matrix from the model. Values are scaled using regulatory norm, with strong links often exceeding 5 or 10.

```python
def get_adj(self):
        """
        Obtain the adjacency matrix. The values in this adjacency matrix have
        been scaled using regulatory norm. You may expect strong links to go
        beyond 5 or 10 in most cases.

        Returns:
            np.ndarray: 2D adjacency matrix of shape (n_gene, n_gene).
        """
        return self.model.get_adj()
```

--------------------------------

### Get GRN Object

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/trainer.html

Obtains a GRN object with the inferred regulatory network. Requires gene names and optionally TF names and a percentile for sparse matrix representation.

```python
def get_grn(self, gene_names, tf_names=None, top_gene_percentile=None):
        """
        Obtain a GRN object. You need to provide the gene names.

        Args:
            gene_names (np.ndarray): An array of names of all genes. The order
                of genes should be the same as the order used in your expression
                table.
            tf_names (np.ndarray): An array of names of all transcriptional
                factors. The order should be the same as the order used in
                your expression table.
            top_gene_percentile (int): If provided, we will set the value on
                weak links to be zero. It is useful if you want to save the
                regulatory relationship in a GRN object as a sparse matrix.

        Returns:
            GRN: A GRN object containing the inferred regulatory network.
        """
        adj = self.model.get_adj()
        return GRN(adj, gene_names, tf_names, top_gene_percentile)
```

--------------------------------

### Finalizing Training and Storing Losses

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/trainer.html

This snippet shows the final steps of the training process, including setting the per-gene losses and calculating the total time cost.

```python
# Set the final per-gene losses, averaged over the batch
self.losses_on_gene = (final_losses_on_gene / batch_size).detach().cpu().numpy()
self.total_time_cost += int((datetime.now() - start_time).total_seconds())
return None
```

--------------------------------

### Moving Tensors Back to CPU

Source: https://tuftsbcb.github.io/RegDiffusion/_autosummary/regdiffusion.models.RegDiffusion.html

Illustrates moving a model's parameters and buffers back to the CPU.

```python
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)
```

--------------------------------

### Initialize RegDiffusionTrainer with Memory Efficiency and AMP

Source: https://tuftsbcb.github.io/RegDiffusion/_sources/large_networks.md

Use this configuration to enable both memory-efficient mode and automatic mixed precision for reduced GPU memory usage. Requires a GPU with bfloat16 support.

```python
trainer = rd.RegDiffusionTrainer(
    exp_array,
    memory_efficient=True,
    use_amp=True,
    device='cuda'
)
trainer.train()
```

--------------------------------

### Discover Target Gene with Strongest Regulation

Source: https://tuftsbcb.github.io/RegDiffusion/_sources/quick_tour.md

Identify the gene with the strongest single regulation from the inferred adjacency matrix. This gene can be used as a starting point for local network analysis.

```python
>>> grn.gene_names[np.argmax(grn.adj_matrix.max(1))]
'HIST1H1D'
```

--------------------------------

### Compute Network Statistics

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/grn.html

Calculates various statistics for the GRN, such as the number of edges, network density, and degree distributions for transcription factors and target genes. Use this to get a quantitative overview of the network's structure.

```python
adj = np.abs(self.adj_matrix)

# Basic statistics
n_edges = np.sum(adj > self.cutoff_threshold)
density = n_edges / (self.n_tfs * self.n_genes)

# Degree statistics
in_degrees = adj.sum(axis=0)
out_degrees = adj.sum(axis=1)

stats = {
    'n_tfs': self.n_tfs,
    'n_genes': self.n_genes,
    'n_edges': int(n_edges),
    'density': float(density),
    'mean_in_degree': float(in_degrees.mean()),
    'std_in_degree': float(in_degrees.std()),
    'max_in_degree': float(in_degrees.max()),
    'mean_out_degree': float(out_degrees.mean()),
    'std_out_degree': float(out_degrees.std()),
    'max_out_degree': float(out_degrees.max()),
    'mean_edge_weight': float(adj[adj > 0].mean()) if n_edges > 0 else 0,
}

return stats
```

--------------------------------

### Training Loop with AMP and Gradient Accumulation

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/trainer.html

This snippet details the core training loop, including forward pass, mixed-precision computation using torch.autocast, MSE loss calculation, gradient accumulation, and optimizer steps. It handles both standard MSE loss and a sparse loss component.

```python
x_noisy, noise = self.forward_pass(x_0, t)
with torch.autocast(
    device_type=self.amp_device_type,
    dtype=torch.bfloat16,
    enabled=self.use_amp
):
    z = self.model(x_noisy, t, ct)

    # Calculate MSE loss, get per-gene loss with reduction='none'
    loss_ = F.mse_loss(noise, z, reduction='none')

# Accumulate per-gene loss values for the last step
with torch.no_grad():
    accumulated_loss_on_gene += loss_.float().squeeze()

# Calculate mean loss for the sample
loss_sample = loss_.mean()
batch_mse_loss += loss_sample.item()

# Scale loss for gradient accumulation
scaled_loss = loss_sample / batch_size
scaled_loss.backward()

# Handle sparse loss once per effective batch
with torch.autocast(
    device_type=self.amp_device_type,
    dtype=torch.bfloat16,
    enabled=self.use_amp
):
    if hasattr(self.model, 'get_sampled_sparse_loss'):
        loss_sparse = self.model.get_sampled_sparse_loss() * self.hp['sparse_loss_coef']
    else:
        adj_m = self.model.get_adj_()
        loss_sparse = adj_m.mean() * self.hp['sparse_loss_coef']

if step_idx > 10:
    loss_sparse.backward()

# Perform the optimizer step after accumulating all gradients
self.opt.step()
```

--------------------------------

### Import Libraries

Source: https://tuftsbcb.github.io/RegDiffusion/_sources/quick_tour.md

Import the necessary regdiffusion and numpy libraries for GRN inference.

```python
import regdiffusion as rd
import numpy as numpy
```

--------------------------------

### Iterating Over All Modules in a Network

Source: https://tuftsbcb.github.io/RegDiffusion/_autosummary/regdiffusion.models.RegDiffusion.html

Demonstrates how to iterate through all modules within a network, including duplicate modules which are yielded only once. Useful for inspecting or modifying all parts of a model.

```python
>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

```

--------------------------------

### load_beeline()

Source: https://tuftsbcb.github.io/RegDiffusion/_autosummary/regdiffusion.data.load_beeline.html

Loads BEELINE data and its ground truth. The data will be downloaded if it doesn't exist locally. It supports multiple benchmark datasets and settings, with specific constraints on data availability for certain settings.

```APIDOC
## load_beeline()

### Description
Loads BEELINE data and its ground truth. The data will be downloaded if it doesn't exist locally. It supports multiple benchmark datasets and settings, with specific constraints on data availability for certain settings.

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Parameters
* **data_dir** (str) - Parent directory to save and load the data. If the path does not exist, it will be created. Data will be saved in a subdirectory under the provided path.
* **benchmark_data** (str) - Benchmark datasets. Choose among “hESC”, “hHep”, “mDC”, “mESC”, “mHSC”, “mHSC-GM”, and “mHSC-L”.
* **benchmark_setting** (str) - Benchmark settings. Choose among “500_STRING”, “1000_STRING”, “500_Non-ChIP”, “1000_Non-ChIP”, “500_ChIP-seq”, “1000_ChIP-seq”, “500_lofgof”, and “1000_lofgof”. If either of the “lofgof” settings is chosen, only “mESC” data is available.

### Returns
A tuple containing two objects for a single BEELINE benchmark. The first element is a scanpy AnnData with cells on rows and genes on columns. Second element is an numpy array for the adjacency list of the ground truth network.

### Return Type
tuple
```

--------------------------------

### RegDiffusionTrainer Initialization

Source: https://tuftsbcb.github.io/RegDiffusion/_autosummary/regdiffusion.RegDiffusionTrainer.html

Initializes the RegDiffusionTrainer with various parameters controlling the model architecture, training process, and device settings.

```APIDOC
## RegDiffusionTrainer

### Description
Initializes and trains a RegDiffusion model. Refer to the paper "From noise to knowledge: probabilistic diffusion-based neural inference" for architecture and training details. The trained model can be accessed via `RegDiffusionTrainer.model`.

### Parameters
* **_exp_array_**: Required. The input data array for training.
* **_cell_types** (Optional): The number of cell types.
* **_T** (int, optional): Diffusion timesteps. Defaults to 5000.
* **_start_noise** (float, optional): Start noise level. Defaults to 0.0001.
* **_end_noise** (float, optional): End noise level. Defaults to 0.02.
* **_time_dim** (int, optional): Dimension for time embedding. Defaults to 64.
* **_celltype_dim** (int, optional): Dimension for cell type embedding. Defaults to 4.
* **_hidden_dims** (list[int], optional): Dimensions of hidden layers in the neural network. Defaults to [16, 16, 16].
* **_init_coef** (int, optional): Initialization coefficient. Defaults to 5.
* **_lr_nn** (float, optional): Learning rate for the neural network. Defaults to 0.001.
* **_lr_adj** (Optional): Learning rate for the adjacency matrix.
* **_weight_decay_nn** (float, optional): Weight decay for the neural network. Defaults to 0.1.
* **_weight_decay_adj** (float, optional): Weight decay for the adjacency matrix. Defaults to 0.01.
* **_sparse_loss_coef** (float, optional): Coefficient for sparse loss. Defaults to 0.25.
* **_adj_dropout** (float, optional): Dropout rate for adjacency matrix. Defaults to 0.3.
* **_batch_size** (int, optional): Batch size for training. Defaults to 128.
* **_n_steps** (int, optional): Number of training steps. Defaults to 1000.
* **_train_split** (float, optional): Fraction of data to use for training. Defaults to 1.0.
* **_train_split_seed** (int, optional): Seed for training split. Defaults to 123.
* **_device** (str, optional): Device to use for training ('cuda' or 'cpu'). Defaults to 'cuda'.
* **_compile** (bool, optional): Whether to compile the model. Defaults to False.
* **_evaluator** (Optional): An evaluator object.
* **_eval_on_n_steps** (int, optional): Evaluate every N steps. Defaults to 100.
* **_logger** (Optional): A logger object.
* **_gradient_accumulation** (bool, optional): Whether to use gradient accumulation. Defaults to False.
* **_use_amp** (bool, optional): Whether to use automatic mixed precision. Defaults to False.
* **_memory_efficient** (bool, optional): Whether to use memory-efficient training. Defaults to False.
* **_seed** (Optional): Random seed for reproducibility.
```

--------------------------------

### Moving Parameters to XPU Device

Source: https://tuftsbcb.github.io/RegDiffusion/_autosummary/regdiffusion.models.RegDiffusion.html

Moves all model parameters and buffers to the XPU, potentially to a specific device.

```python
xpu(_device : int | device | None = None_)
```

--------------------------------

### Power Beta Schedule for Diffusion

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/trainer.html

Generates a power-law beta schedule for diffusion models. This allows for non-linear control over noise levels.

```python
import torch

def power_beta_schedule(timesteps, start_noise, end_noise, power=2):
    linspace = torch.linspace(0, 1, timesteps, dtype = torch.float)
    poweredspace = linspace ** power
    scale = 1000 / timesteps
    beta_start = scale * start_noise
    beta_end = scale * end_noise
    return beta_start + (beta_end - beta_start) * poweredspace
```

--------------------------------

### Training with Gradient Accumulation

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/trainer.html

This snippet demonstrates training using gradient accumulation. Gradients are computed for each sample individually and accumulated before an optimizer step, which reduces memory usage at the cost of increased training time.

```python
        start_time = datetime.now()
        eval_steps = self.hp['eval_on_n_steps']
        if n_steps is None:
            n_steps = self.hp['n_steps']
        
        batch_size = self.hp['batch_size']
        sampled_adj = self.model.get_sampled_adj_()
        
        # This will hold the per-gene losses from the final step
        final_losses_on_gene = torch.zeros(self.n_gene, device=self.device)

        with tqdm(range(n_steps)) as pbar:
            for step_idx in pbar:
                if self.hp['seed'] is not None:
                    torch.manual_seed(self.hp['seed'] + step_idx)
                # Fetch a single batch for the current step
                x_0_batch, ct_batch = next(iter(self.train_dataloader))
                x_0_batch = x_0_batch.to(self.device)
                ct_batch = ct_batch.to(self.device)
                
                # Zero gradients before starting the accumulation for the batch
                self.opt.zero_grad() 
                
                batch_mse_loss = 0.0
                accumulated_loss_on_gene = torch.zeros(self.n_gene, device=self.device)

                # Loop over each sample in the batch to accumulate gradients
                for i in range(batch_size):
                    x_0 = x_0_batch[i:i+1]
                    ct = ct_batch[i:i+1]
                    
                    # Generate timestep for the single sample
                    t = torch.randint(0, self.hp['T'], (1,), device=self.device).long()

```

--------------------------------

### RegDiffusion Model Initialization

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/models/regdiffusion.html

Initializes the RegDiffusion model, setting up parameters for gene count, dimensions, dropout, and the adjacency matrix.

```python
import torch
from torch import nn

class RegDiffusion(nn.Module):
    def __init__(self,
        n_gene, time_dim, 
        n_celltype=None, celltype_dim=4, 
        hidden_dims=[16, 16, 16], adj_dropout=0.3, init_coef = 5
    ):
        super(RegDiffusion, self).__init__()
        
        self.n_gene = n_gene
        self.gene_dim = hidden_dims[0]
        self.adj_dropout=adj_dropout
        self.gene_reg_norm = 1/(n_gene-1)
        self.celltype_dim = celltype_dim
        self.n_celltype = n_celltype
        
        adj_A = torch.ones(n_gene, n_gene) * self.gene_reg_norm * init_coef
        self.adj_A = nn.Parameter(adj_A, requires_grad =True, )
        self.sampled_adj_row_nonparam = nn.Parameter(
            torch.randint(0, n_gene, size=(10000,)), 
            requires_grad=False)
        self.sampled_adj_col_nonparam = nn.Parameter(
            torch.randint(0, n_gene, size=(10000,)),
            requires_grad=False)
```

--------------------------------

### GRNEvaluator Class Initialization

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/evaluator.html

Initializes the GRNEvaluator with ground truth, gene names, and desired metrics. Ensure gene names match the order in the adjacency matrix.

```python
import numpy as np
import pandas as pd
from sklearn.metrics import average_precision_score, roc_auc_score



[docs]
class GRNEvaluator:
    """
    A generalized evaluator for GRN inference.

    Args:
        ground_truth (np.ndarray or list): Either a 2D numpy array or list of 
            list holding the ground truth. Each row is an edge and includes 
            names for the source and target nodes. For example, [['A', 'B'], 
            ['B', 'C']].
        gene_names (np.ndarray or list): Either a 1D numpy array or list of 
            gene names. Make sure the order of the gene names is the same as 
            the order of gene names in the adjacency matrix. 
        metrics (list): A list of supported evaluation metrics. Currently 
            support 'AUROC', 'AUPR', 'AUPRR', 'EP', 'EPR'.
    """
    def __init__(self, ground_truth, gene_names, 
                 metrics=['AUROC', 'AUPR', 'AUPRR', 'EP', 'EPR']):
        n_gene = len(gene_names)
        gene1 = [x[0] for x in ground_truth]
        gene2 = [x[1] for x in ground_truth]
        
        # TF is the set of genes appearing as regulators
        TF = set(gene1)
        # All_gene is the combined set of regulator and target genes
        All_gene = set(gene1) | set(gene2)
    
        tf_mask = (np.zeros(n_gene) == 1)
        # Not all provided genes are in the all_gene set
        gene_mask = (np.zeros(n_gene) == 1)
        tf_map = {}
        gene_map = {}
        for i, item in enumerate(gene_names):
            if item in TF:
                tf_mask[i] = True
                tf_map[item] = len(tf_map)
            if item in All_gene:
                gene_mask[i] = True
                gene_map[item] = len(gene_map)
        
        y_true = np.zeros([len(TF), len(All_gene)])
        for link in ground_truth:
            y_true[tf_map[link[0]], gene_map[link[1]]] = 1.0
        y_true = y_true.flatten()

        self.tf_mask = tf_mask
        self.gene_mask = gene_mask
        self.y_true = y_true
        self.y_true_mean = np.mean(y_true)
        self.num_true_edges = int(y_true.sum())
        self.ground_truth = ground_truth
        self.report_auroc = ('AUROC' in metrics)
        self.report_aupr = ('AUPR' in metrics)
        self.report_auprr = ('AUPRR' in metrics)
        self.report_ep = ('EP' in metrics)
        self.report_epr = ('EPR' in metrics)

```

--------------------------------

### Enable Mixed Precision Training (AMP)

Source: https://tuftsbcb.github.io/RegDiffusion/large_networks.html

Combine memory-efficient mode with Automatic Mixed Precision (AMP) for a modest additional memory reduction. This requires a GPU with bfloat16 support and may increase training time.

```python
trainer = rd.RegDiffusionTrainer(
    exp_array,
    memory_efficient=True,
    use_amp=True,
    device='cuda'
)
trainer.train()

```

--------------------------------

### GRNEvaluator Constructor

Source: https://tuftsbcb.github.io/RegDiffusion/_autosummary/regdiffusion.GRNEvaluator.html

Initializes the GRNEvaluator with ground truth, gene names, and a list of metrics.

```APIDOC
## GRNEvaluator

### Description
A generalized evaluator for GRN inference.

### Parameters

*   **ground_truth** (np.ndarray or list) - Required - Either a 2D numpy array or list of list holding the ground truth. Each row is an edge and includes names for the source and target nodes. For example, [['A', 'B'], ['B', 'C']].
*   **gene_names** (np.ndarray or list) - Required - Either a 1D numpy array or list of gene names. Make sure the order of the gene names is the same as the order of gene names in the adjacency matrix.
*   **metrics** (list) - Optional - A list of supported evaluation metrics. Currently support ‘AUROC’, ‘AUPR’, ‘AUPRR’, ‘EP’, ‘EPR’. Default is ['AUROC', 'AUPR', 'AUPRR', 'EP', 'EPR'].
```

--------------------------------

### Standard Training Loop with Evaluation

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/trainer.html

This snippet illustrates a standard training loop that includes periodic evaluation and logging. It calculates training loss, tracks changes in the adjacency matrix, and computes validation loss if a validation dataloader is provided.

```python
                    if epoch > 10:
                        loss = loss + loss_sparse
                    loss.backward()
                    self.opt.step()
                    epoch_loss.append(loss.item())
                train_loss = np.mean(epoch_loss)
                sampled_adj_new = self.model.get_sampled_adj_()
                adj_diff = (
                    sampled_adj_new - sampled_adj
                    ).mean().item()*(self.n_gene-1)
                sampled_adj = sampled_adj_new
                pbar.set_description(
                    f'Training loss: {train_loss:.3f}, Change on Adj: {adj_diff:.3f}')
                epoch_log = {'train_loss': train_loss, 'adj_change': adj_diff}
                if epoch % eval_steps == eval_steps - 1:
                    if self.evaluator is not None:
                        eval_result = self.evaluator.evaluate(
                            self.model.get_adj()
                            )
                        for k in eval_result.keys():
                            epoch_log[k] = eval_result[k]
                    if self.hp['train_split'] < 1:
                        if self.hp['seed'] is not None:
                            torch.manual_seed(self.hp['seed'] + n_steps + epoch)
                        with torch.no_grad():
                            val_epoch_loss = []
                            for step, batch in enumerate(self.val_dataloader):
                                x_0, ct = batch
                                x_0 = x_0.to(self.device)
                                ct = ct.to(self.device)
                                t = torch.randint(
                                    0, self.hp['T'], (x_0.shape[0],),
                                    device=self.device).long()

                                x_noisy, noise = self.forward_pass(x_0, t)
                                with torch.autocast(
                                    device_type=self.amp_device_type,
                                    dtype=torch.bfloat16,
                                    enabled=self.use_amp
                                ):
                                    z = self.model(x_noisy, t, ct)
                                    step_val_loss = F.mse_loss(
                                        noise, z, reduction='mean').item()
                                val_epoch_loss.append(step_val_loss)
                            epoch_log['val_loss'] = np.mean(val_epoch_loss)
                self.logger.log(epoch_log)
        self.losses_on_gene = loss_.detach().mean(0).cpu().numpy()
        self.total_time_cost += int(
            (datetime.now() - start_time).total_seconds())
        return None
```

--------------------------------

### Train Model (Normal Mode)

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/trainer.html

Trains the initialized model for a specified number of steps. Handles seeding and loss calculation, including optional sparse loss.

```python
start_time = datetime.now()
        eval_steps = self.hp['eval_on_n_steps']
        if n_steps is None:
            n_steps = self.hp['n_steps']
        sampled_adj = self.model.get_sampled_adj_()
        with tqdm(range(n_steps)) as pbar:
            for epoch in pbar:
                if self.hp['seed'] is not None:
                    torch.manual_seed(self.hp['seed'] + epoch)
                epoch_loss = []
                for step, batch in enumerate(self.train_dataloader):
                    x_0, ct = batch
                    x_0 = x_0.to(self.device)
                    ct = ct.to(self.device)
                    self.opt.zero_grad()
                    t = torch.randint(
                        0, self.hp['T'], (x_0.shape[0],),
                        device=self.device
                    ).long()
        
                    x_noisy, noise = self.forward_pass(x_0, t)
                    with torch.autocast(
                        device_type=self.amp_device_type,
                        dtype=torch.bfloat16,
                        enabled=self.use_amp
                    ):
                        z = self.model(x_noisy, t, ct)
                        loss_ = F.mse_loss(noise, z, reduction='none')
                        loss = loss_.mean()

                        if hasattr(self.model, 'get_sampled_sparse_loss'):
                            loss_sparse = self.model.get_sampled_sparse_loss() * self.hp['sparse_loss_coef']
                        else:
                            adj_m = self.model.get_adj_()
                            loss_sparse = adj_m.mean() * self.hp['sparse_loss_coef']
```

--------------------------------

### Moving Parameters to Empty Device

Source: https://tuftsbcb.github.io/RegDiffusion/_autosummary/regdiffusion.models.RegDiffusion.html

Moves model parameters and buffers to a specified device without copying storage.

```python
to_empty(_*_ , _device : str | device | int | None_, _recurse : bool = True_)
```

--------------------------------

### Iterating Through Named Modules

Source: https://tuftsbcb.github.io/RegDiffusion/_autosummary/regdiffusion.models.RegDiffusion.html

Demonstrates how to iterate through all named modules within a sequential model. Duplicate modules are returned only once.

```python
>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
```

--------------------------------

### Train RegDiffusion Model

Source: https://tuftsbcb.github.io/RegDiffusion/_sources/downstream_with_pyscenic.md

Initializes and trains the RegDiffusion model using the prepared log-transformed count data. Training time varies based on hardware.

```python
rd_trainer = rd.RegDiffusionTrainer(x)
rd_trainer.train()
```

--------------------------------

### Visualizing UMAP with Sample Labels and Batch Effect Comparison

Source: https://tuftsbcb.github.io/RegDiffusion/downstream_with_pyscenic.html

Visualizes UMAP results colored by sample labels to assess batch effects, comparing PCA, Harmony-corrected PCA, and AUCell-based UMAPs.

```python
label_dict = {
    "SRR30599518": "CTL-1",
    "SRR30599519": "CTL-2",
    "SRR30599520": "CTL-3",
    "SRR30599521": "LPS-1",
    "SRR30599522": "LPS-2",
    "SRR30599523": "LPS-3"
}
adata.obs['label'] = [label_dict[x] for x in adata.obs['sample']]
fig, axs = plt.subplots(1, 3, figsize=(13, 4))

sc.pl.scatter( adata, basis='pca_umap', 
    color=['label'],
    title=['PCA - UMAP'],
    alpha=0.9, ax=axs[0], show=False, palette='RdYlBu'
    )

sc.pl.scatter( adata, basis='harmony_pca_umap', 
    color=['label'],
    title=['Harmony - PCA - UMAP'],
    alpha=0.9, ax=axs[1], show=False, palette='RdYlBu'
    )

sc.pl.scatter( adata, basis='aucell_umap', 
    color=['label'],
    title=['RegDiffusion - AUCell - UMAP'],
    alpha=0.9, ax=axs[2], show=False, palette='RdYlBu'
    )

axs[0].get_legend().remove()
axs[1].get_legend().remove()

plt.tight_layout()
plt.show()
```

--------------------------------

### Prepare Training and Validation DataLoaders

Source: https://tuftsbcb.github.io/RegDiffusion/_modules/regdiffusion/trainer.html

Splits normalized expression data and cell types into training and validation sets. Creates PyTorch TensorDatasets and DataLoaders for both training and validation, with random sampling for the training dataloader.

```python
        random_state = np.random.RandomState(train_split_seed)
        train_val_split = random_state.rand(normalized_X.shape[0])
        train_index = train_val_split <= train_split
        val_index = train_val_split > train_split

        x_tensor_train = torch.tensor(
            normalized_X[train_index, ], dtype=torch.float32)
        celltype_tensor_train = torch.tensor(
            cell_types[train_index], dtype=int)
        x_tensor_val = torch.tensor(
            normalized_X[val_index, ], dtype=torch.float32)
        celltype_tensor_val = torch.tensor(cell_types[val_index], dtype=int)

        self.train_dataset = TensorDataset(
            x_tensor_train, celltype_tensor_train)
        train_sampler = torch.utils.data.RandomSampler(
            self.train_dataset, replacement=True, num_samples=batch_size,
            generator=self._sampler_generator)
        self.train_dataloader = DataLoader(
            self.train_dataset,
            sampler=train_sampler,
            batch_size=batch_size,
            drop_last=True)

        self.val_dataset = TensorDataset(
            x_tensor_val, celltype_tensor_val)
        self.val_dataloader = DataLoader(
            self.val_dataset,
            shuffle=False,
            batch_size=batch_size,
            drop_last=False)
```