dspin 1.3.4__tar.gz → 1.4.0__tar.gz

{dspin-1.3.4/dspin.egg-info → dspin-1.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: dspin
-Version: 1.3.4
+Version: 1.4.0
 Summary: Regulatory network models from single-cell perturbation profiling
 Home-page: https://github.com/JialongJiang/DSPIN
 Author: Jialong Jiang
@@ -25,6 +25,18 @@ Requires-Dist: scanpy
 Requires-Dist: tqdm
 Requires-Dist: igraph
 Requires-Dist: leidenalg
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
 
 # D-SPIN: Constructing regulatory network models from scRNA-seq perturbation profiling  
 

dspin-1.4.0/dspin/__init__.py

@@ -0,0 +1,18 @@
+
+from . import annotate
+from . import compute
+from . import plot
+from . import dspin
+
+# Import main classes for easy access
+from .dspin import DSPIN, GeneDSPIN, ProgramDSPIN, AbstractDSPIN
+
+__all__ = [
+    'DSPIN',
+    'GeneDSPIN', 
+    'ProgramDSPIN',
+    'AbstractDSPIN',
+    'annotate',
+    'compute',
+    'plot'
+]

{dspin-1.3.4 → dspin-1.4.0}/dspin/compute.py

@@ -232,15 +232,16 @@ def onmf(X: np.array, rank: int, max_iter: int = 500) -> Tuple[np.array, np.arra
         coef_S = AtX / S.dot(AtX.T).dot(S)
         S = np.nan_to_num(S * coef_S) #, posinf=1e5)
 
-        pbar.update(1)
+        # pbar.update(1)
 
         # Calculate the reconstruction error every 10 iterations and update the
         # progress bar
-        if itr % 10 == 0:
+        if itr % 20 == 0:
             error = np.linalg.norm(X - np.dot(A, S), 'fro')
             self_product = S.dot(S.T)
             orthogonal_error = np.linalg.norm(self_product - np.diag(np.diag(self_product)), 'fro') / np.linalg.norm(S, 'fro')
             pbar.set_postfix({"Reconstruction error": f"{error:.2f}", "Orthogonal error": f"{orthogonal_error:.4f}"})
+            pbar.update(20)
 
     pbar.close()
 
@@ -278,6 +279,8 @@ def compute_onmf(seed: int,
     """
     # Generate a random seed
     np.random.seed(seed)
+
+    tqdm._instances.clear()
     # Factorized Matrices
     H, W = onmf(gene_matrix, num_spin, max_iter)
 
@@ -608,7 +611,7 @@ def para_moments(j_mat: np.ndarray,
             spins[s] = (tmp % 3) - 1
             tmp //= 3
 
-        # Energy: E = −h·s − ½ sᵀJs
+        # Energy: E = −h·s − s'Js / 2
         energy: float = 0.0
         for ii in range(num_spin):
             energy -= h_vec[ii] * spins[ii]
@@ -796,6 +799,28 @@ def samp_moments(j_mat: np.ndarray,
     return corr_para, mean_para
 
 
+def _gradient_one_round(kk: int,
+                        cur_j: np.ndarray,
+                        cur_h: np.ndarray,
+                        raw_data_k,
+                        method: str,
+                        train_dat: Dict[str, Any]):
+    """Helper that compute gradient for a single round."""
+    if method == 'pseudo_likelihood':
+        j_grad, h_grad = pseudol_gradient(cur_j, cur_h[:, kk:kk + 1], raw_data_k, directed=train_dat['directed'])
+        h_grad = h_grad.flatten()
+    else:
+        if method == 'maximum_likelihood':
+            corr_para, mean_para = para_moments(cur_j, cur_h[:, kk])
+        elif method == 'mcmc_maximum_likelihood':
+            corr_para, mean_para = samp_moments(
+                    cur_j, cur_h[:, kk], train_dat['mcmc_samplingsz'], train_dat['mcmc_samplingmix'], train_dat['mcmc_samplegap'])
+        j_grad  = corr_para - raw_data_k[0]
+        h_grad  = mean_para  - raw_data_k[1].flatten()
+
+    return kk, j_grad, h_grad
+
+
 def compute_gradient(cur_j: np.ndarray,
                             cur_h: np.ndarray,
                             raw_data: List[Tuple[np.ndarray, np.ndarray]],
@@ -828,49 +853,32 @@ def compute_gradient(cur_j: np.ndarray,
     rec_jgrad = np.zeros((num_spin, num_spin, num_round))
     rec_hgrad = np.zeros((num_spin, num_round))
 
-    for kk in range(num_round):
-        # Computing gradients using the specified method.
-        if method == 'pseudo_likelihood':
+    if method == 'pseudo_likelihood':
+         for kk in range(num_round):
+            # Computing gradients using the specified method.
             j_grad, h_grad = pseudol_gradient(
                 cur_j, cur_h[:, kk: kk + 1], raw_data[kk], directed=train_dat['directed'])
             h_grad = h_grad.flatten()
-        else:
-            # Distinguishing between other methods and computing gradients
-            # accordingly
-            if method == 'maximum_likelihood':
-                corr_para, mean_para = para_moments(cur_j, cur_h[:, kk])
-            elif method == 'mcmc_maximum_likelihood':
-                corr_para, mean_para = samp_moments(
-                    cur_j, cur_h[:, kk], train_dat['mcmc_samplingsz'], train_dat['mcmc_samplingmix'], train_dat['mcmc_samplegap'])
-            j_grad = corr_para - raw_data[kk][0]
-            h_grad = mean_para - raw_data[kk][1].flatten()
-
-        rec_jgrad[:, :, kk] = j_grad
-        rec_hgrad[:, kk] = h_grad
-
-    return rec_jgrad, rec_hgrad
 
-
-def _gradient_one_round(kk: int,
-                        cur_j: np.ndarray,
-                        cur_h: np.ndarray,
-                        raw_data_k,
-                        method: str,
-                        train_dat: Dict[str, Any]):
-    """Helper that compute gradient for a single round."""
-    if method == 'pseudo_likelihood':
-        j_grad, h_grad = pseudol_gradient(cur_j, cur_h[:, kk:kk + 1], raw_data_k, directed=train_dat['directed'])
-        h_grad = h_grad.flatten()
+            rec_jgrad[:, :, kk] = j_grad
+            rec_hgrad[:, kk] = h_grad
     else:
-        if method == 'maximum_likelihood':
-            corr_para, mean_para = para_moments(cur_j, cur_h[:, kk])
-        elif method == 'mcmc_maximum_likelihood':
-            corr_para, mean_para = samp_moments(
-                    cur_j, cur_h[:, kk], train_dat['mcmc_samplingsz'], train_dat['mcmc_samplingmix'], train_dat['mcmc_samplegap'])
-        j_grad  = corr_para - raw_data_k[0]
-        h_grad  = mean_para  - raw_data_k[1].flatten()
-
-    return kk, j_grad, h_grad
+        n_jobs = train_dat.get('max_workers', min(num_round, os.cpu_count())) 
+        chunk = train_dat.get('chunk_size', 1)          # rounds per task
+        backend = 'loky'                                # process‑based
+
+        # print(f'Processing with {n_jobs} workers...')
+        with parallel_backend(backend, n_jobs=n_jobs):
+            results = Parallel(batch_size=chunk)(
+                delayed(_gradient_one_round)(kk, cur_j, cur_h, raw_data[kk], method, train_dat)
+                for kk in range(num_round)
+            )
+
+        for kk, j_grad, h_grad in results:
+            rec_jgrad[:, :, kk] = j_grad
+            rec_hgrad[:, kk] = h_grad
+        
+    return rec_jgrad, rec_hgrad
 
 
 def compute_gradient_para(cur_j: np.ndarray,
@@ -990,10 +998,13 @@ def apply_regularization(rec_jgrad: np.ndarray,
     
     if 'lambda_l2_j_prior_mask' in train_dat:
         rec_jgrad += train_dat['lambda_l2_j_prior_mask'] * (cur_j * (train_dat['j_prior_mask'] == 0))[:, :, np.newaxis]
+
+    if 'lambda_l1_j_prior_mask' in train_dat:
+        rec_jgrad += train_dat['lambda_l1_j_prior_mask'] * ((cur_j * (train_dat['j_prior_mask'] == 0))[:, :, np.newaxis] / 1e-3).clip(-1, 1)
     
     if print_regu:  
         print('Regularization parameters')
-        all_regu_list = ['lambda_l1_j', 'lambda_l1_h', 'lambda_l2_j', 'lambda_l2_h', 'lambda_l2_j_prior', 'lambda_l2_j_prior_mask', 'lambda_l2_h_rela_prior', 'lambda_l1_h_rela', 'lambda_l2_h_rela']
+        all_regu_list = ['lambda_l1_j', 'lambda_l1_h', 'lambda_l2_j', 'lambda_l2_h', 'lambda_l2_j_prior', 'lambda_l2_j_prior_mask', 'lambda_l2_h_rela_prior', 'lambda_l1_h_rela', 'lambda_l2_h_rela', 'lambda_l1_j_prior_mask']
         all_values = [train_dat.get(key, 0) for key in all_regu_list]
         for key, value in zip(all_regu_list, all_values):
             if value > 0:
@@ -1153,7 +1164,6 @@ def learn_network_adam(raw_data: Any,
         if (counter > backtrack_gap) and (rec_jgrad_sum_norm[counter - 1] > 1.5 * rec_jgrad_sum_norm[counter - 1 - backtrack_gap]):
             if_backtrack = True
             
-
         if (counter > 2 * backtrack_gap) and (counter % backtrack_gap == 0):
 
             est_cur_grad = rec_jgrad_sum_norm[counter - 1 - backtrack_gap: counter - 1].mean()
@@ -1162,7 +1172,6 @@ def learn_network_adam(raw_data: Any,
             if est_cur_grad > 1.05 * est_prev_grad:
                 if_backtrack = True
 
-
         if if_backtrack:
             print('Backtracking at epoch %d' % counter)
             backtrack_counter += 1
@@ -1180,8 +1189,7 @@ def learn_network_adam(raw_data: Any,
 
     pbar.close()
 
-    # Retrieve parameters corresponding to the minimum gradient norm found
-    # during training
+    # Retrieve parameters corresponding to the minimum gradient norm found during training
     trace_epoch = 50 + num_epoch - counter
     pos = num_epoch - trace_epoch + np.argmin(rec_jgrad_sum_norm[- trace_epoch: ])
     # print(pos)
@@ -1221,16 +1229,14 @@ def learn_program_regulators(gene_states: List[np.ndarray],
     num_round = len(gene_states)
 
     num_epoch, stepsz, rec_gap = (
-        train_dat.get(
-            key, None) for key in [
+        train_dat.get(key, None) for key in [
             "num_epoch", "stepsz", "rec_gap"])
     list_step = np.arange(num_epoch, 0, - rec_gap)[::-1]
 
     save_path = train_dat.get('save_path', None)
 
     backtrack_gap, backtrack_tol = (
-        train_dat.get(
-            key, None) for key in [
+        train_dat.get(key, None) for key in [
             "backtrack_gap", "backtrack_tol"])
 
     lambda_l1_j, lambda_l2_j = (train_dat.get(key, 0) for key in ["lambda_l1_j", "lambda_l2_j"])
@@ -1246,6 +1252,10 @@ def learn_program_regulators(gene_states: List[np.ndarray],
     samp_weight = samp_weight / np.sum(samp_weight)
     
     # Preallocate gradient norm recording (optional)
+    rec_interaction_all = np.zeros((num_epoch, num_gene, num_program))
+    rec_selfj_all = np.zeros((num_epoch, num_program))
+    rec_selfh_all = np.zeros((num_epoch, num_program))
+    
     rec_grad = np.zeros(num_epoch)
 
     m_interaction, v_interaction = np.zeros_like(cur_interaction), np.zeros_like(cur_interaction)
@@ -1265,9 +1275,14 @@ def learn_program_regulators(gene_states: List[np.ndarray],
     rec_interaction_grad = np.zeros((num_round, num_gene, num_program))
     rec_selfj_grad = np.zeros((num_round, num_program))
     rec_selfh_grad = np.zeros((num_round, num_program))
-    
+
+    log_adam_grad = {name: deque(maxlen=backtrack_gap)
+                     for name in ["m_interaction", "v_interaction", "m_selfj", "v_selfj", "m_selfh", "v_selfh"]}
+
+    tqdm._instances.clear()
     pbar = tqdm(total=num_epoch)
     counter = 1
+    backtrack_counter = 0
 
     while counter <= num_epoch:
 
@@ -1320,15 +1335,52 @@ def learn_program_regulators(gene_states: List[np.ndarray],
         update_val, m_selfh, v_selfh = update_adam(selfh_grad, m_selfh, v_selfh, counter, stepsz)
         cur_selfh = cur_selfh - update_val
         
+        rec_interaction_all[counter - 1, :, :] = cur_interaction
+        rec_selfj_all[counter - 1, :] = cur_selfj
+        rec_selfh_all[counter - 1, :] = cur_selfh
+        
+        for name, value in zip(["m_interaction", "v_interaction", "m_selfj", "v_selfj", "m_selfh", "v_selfh"], [m_interaction, v_interaction, m_selfj, v_selfj, m_selfh, v_selfh]):
+            log_adam_grad[name].append(value)
+        
         # Record gradient norm (optional)
         rec_grad[counter - 1] = np.linalg.norm(interaction_grad)
         
         if counter in list_step:
             pbar.update(rec_gap)
-            pbar.set_postfix({"Network Gradient": f"{rec_grad[counter - 1]:.4f}"})
+            pbar.set_postfix({"Network Gradient": f"{rec_grad[counter - 1]:.6f}"})
+
+        if_backtrack = False
+        if (counter > backtrack_gap) and (rec_grad[counter - 1] > 1.5 * rec_grad[counter - 1 - backtrack_gap]):
+            if_backtrack = True
+            
+        if (counter > 2 * backtrack_gap) and (counter % backtrack_gap == 0):
+            est_cur_grad = rec_grad[counter - 1 - backtrack_gap: counter - 1].mean()
+            est_prev_grad = rec_grad[counter - 1 - 2 * backtrack_gap: counter - 1 - backtrack_gap].mean()
+            
+            if est_cur_grad > 1.05 * est_prev_grad:
+                if_backtrack = True
 
-        counter += 1
+        if if_backtrack:
+            print('Backtracking at epoch %d' % counter)
+            backtrack_counter += 1
+            m_interaction, v_interaction, m_selfj, v_selfj, m_selfh, v_selfh = [log_adam_grad[key][0] for key in ["m_interaction", "v_interaction", "m_selfj", "v_selfj", "m_selfh", "v_selfh"]]
+            counter = counter - backtrack_gap
+            stepsz = stepsz / 4
+            if backtrack_counter > backtrack_tol:
+                print(f'Backtracking more than {backtrack_tol} times, stop training.')
+                break
+        else:   
+            counter += 1
     
+    pbar.close()
+
+    # Retrieve parameters corresponding to the minimum gradient norm found during training
+    trace_epoch = 50 + num_epoch - counter
+    pos = num_epoch - trace_epoch + np.argmin(rec_grad[- trace_epoch: ])
+    cur_interaction = rec_interaction_all[pos, :, :]
+    cur_selfj = rec_selfj_all[pos, :]
+    cur_selfh = rec_selfh_all[pos, :]
+
     return cur_interaction, cur_selfj, cur_selfh
 
 

{dspin-1.3.4 → dspin-1.4.0}/dspin/dspin.py

@@ -35,6 +35,9 @@ import warnings
 import itertools
 from typing import List, Dict, Optional, Tuple, Union
 from scipy.stats import skew
+import leidenalg as la
+import igraph as ig
+import matplotlib as mpl
 
 
 class AbstractDSPIN(ABC):
@@ -107,6 +110,18 @@ class AbstractDSPIN(ABC):
     def sample_list(self):
         return self._samp_list
 
+    @property
+    def modules(self):
+        return self._modules
+    
+    @property
+    def name_list(self):
+        return self._name_list
+
+    @property
+    def name_list_short(self):
+        return self._name_list_short
+
     @program_representation_raw.setter
     def program_representation_raw(self, value):
         self._onmf_rep_ori = value
@@ -131,6 +146,18 @@ class AbstractDSPIN(ABC):
     def sample_list(self, value):
         self._samp_list = value
 
+    @modules.setter
+    def modules(self, value):
+        self._modules = value
+
+    @name_list.setter
+    def name_list(self, value):
+        self._name_list = value
+
+    @name_list_short.setter
+    def name_list_short(self, value):
+        self._name_list_short = value
+
     def discretize(self, 
                    clip_percentile: float = 100,
                    num_init: int = 10) -> None:
@@ -201,11 +228,11 @@ class AbstractDSPIN(ABC):
         num_spin = self.num_spin
         raw_data = self._raw_data
 
-        if self.example_list is not None:
-            example_list_ind = [list(self._samp_list).index(samp)
-                                for samp in self.example_list]
-            self._samp_list = self._samp_list[example_list_ind]
-            raw_data = raw_data[example_list_ind]
+        if self.sample_list_ordered is not None:
+            sample_list_ordered_ind = [list(self._samp_list).index(samp)
+                                for samp in self.sample_list_ordered]
+            self._samp_list = self._samp_list[sample_list_ordered_ind]
+            raw_data = raw_data[sample_list_ordered_ind]
             self._raw_data = raw_data
             print("Using a subset of {} samples for network inference.".format(len(self._samp_list)))
 
@@ -241,7 +268,9 @@ class AbstractDSPIN(ABC):
                           method: str = 'auto',
                           directed: bool = False,
                           params: dict = None,
-                          example_list: List[str] = None,
+                          sample_list_ordered: List[str] = None,
+                          prior_network: np.ndarray = None,
+                          perturb_matrix: np.ndarray = None,
                           run_with_matlab: bool = False) -> None:
         """
         Perform network inference using a specified method and parameters.
@@ -256,8 +285,12 @@ class AbstractDSPIN(ABC):
             Whether to infer a directed network. Default is False.
         params : dict, optional
             Additional parameters for network inference. Default is None.
-        example_list : List[str], optional
-            List of example sample identifiers. Default is None.
+        sample_list_ordered : List[str], optional
+            List of ordered sample identifiers. Default is None.
+        prior_network : np.ndarray, optional
+            Binary matrix of preferred network connectivity from prior knowledge. Default is None.
+        perturb_matrix : np.ndarray, optional
+            Matrix of putative relative responses to control samples. Default is None.
         run_with_matlab : bool, optional
             If True, prepares data for MATLAB execution instead of running inference in Python. Default is False.
         """
@@ -275,8 +308,8 @@ class AbstractDSPIN(ABC):
         if method == 'auto':           
             samp_list = np.unique(self.adata.obs[sample_id_key])
             num_sample = len(samp_list)
-            if example_list is not None:
-                num_sample = len(example_list)
+            if sample_list_ordered is not None:
+                num_sample = len(sample_list_ordered)
             if num_sample > 30:
                 method = 'pseudo_likelihood'
             else:
@@ -296,9 +329,18 @@ class AbstractDSPIN(ABC):
         else:
             self.raw_data_corr(sample_id_key)
 
-        self.example_list = example_list
+        self.sample_list_ordered = sample_list_ordered
 
         train_dat = self.default_params(method)
+
+        if prior_network is not None:
+            train_dat['j_prior_mask'] = prior_network
+            train_dat['lambda_l1_j_prior_mask'] = 0.01
+
+        if perturb_matrix is not None:
+            train_dat['perturb_matrix'] = perturb_matrix
+            train_dat['lambda_l2_h_rela_prior'] = 0.5
+
         train_dat['directed'] = directed
         if params is not None:
             train_dat.update(params)
@@ -347,6 +389,62 @@ class AbstractDSPIN(ABC):
 
         self._relative_responses = compute_relative_responses(self.responses, if_control, batch_index)
 
+    def compute_modules(self, 
+                        resolution: float = 1.0,
+                        seed: int = 0,
+                        thres: float = 0) -> List[List[int]]:
+        """
+        Compute network modules using Leiden community detection.
+        
+        Parameters
+        ----------
+        resolution : float, optional
+            Resolution parameter for Leiden community detection. Higher values lead to more modules. Default is 1.0.
+        seed : int, optional
+            Random seed for Leiden community detection. Default is 0.
+        thres : float, optional
+            Threshold for filtering weak connections. Default is 0.
+            
+        Returns
+        -------
+        List[List[int]]
+            List of modules, where each module is a list of node indices.
+        """
+        if not hasattr(self, '_network') or self._network is None:
+            raise ValueError("Network must be inferred before computing modules. Run network_inference() first.")
+            
+        j_mat = self._network.copy()
+        np.fill_diagonal(j_mat, 0)
+        
+        # Filter weak connections
+        j_filt = j_mat.copy()
+        j_filt[np.abs(j_mat) < thres] = 0
+        np.fill_diagonal(j_filt, 0)
+        
+        # Create networkx graph
+        G = nx.from_numpy_array(j_filt)
+        G = ig.Graph.from_networkx(G)
+        
+        # Separate positive and negative edges
+        G_pos = G.subgraph_edges(G.es.select(weight_gt=0), delete_vertices=False)
+        G_neg = G.subgraph_edges(G.es.select(weight_lt=0), delete_vertices=False)
+        G_neg.es['weight'] = [-w for w in G_neg.es['weight']]
+        
+        # Apply Leiden community detection
+        part_pos = la.RBConfigurationVertexPartition(G_pos, weights='weight', resolution_parameter=resolution)
+        part_neg = la.RBConfigurationVertexPartition(G_neg, weights='weight', resolution_parameter=resolution)
+        
+        optimiser = la.Optimiser()
+        optimiser.set_rng_seed(seed)
+        
+        diff = optimiser.optimise_partition_multiplex([part_pos, part_neg], layer_weights=[1, -1])
+        
+        # Extract modules
+        net_class = list(part_pos)
+        self._modules = net_class
+        
+        print(f'Identified {len(net_class)} modules with sizes {np.array([len(cur_list) for cur_list in net_class])}')
+
 
 class GeneDSPIN(AbstractDSPIN):
     """
@@ -392,7 +490,6 @@ class GeneDSPIN(AbstractDSPIN):
 
         if gene_skew.mean() > 2:
             default_clip = 95
-            print("Input gene expression seem to be heavy-tailed. Using 95 percentile for clipping before discretization.")
         else:
             default_clip = 100
 
@@ -400,13 +497,20 @@ class GeneDSPIN(AbstractDSPIN):
                    'clip_percentile': default_clip,
                    'num_init': 10}
         dparams.update(discretize_params)
-        
+
+        self.name_list = self.adata.var_names.tolist()
+        self.name_list_short = self.adata.var_names.tolist()
+
         if dparams['use_default_discretize']:
+
+            if dparams['clip_percentile'] < 100:
+                print(f"Input gene expression seem to be heavy-tailed. Using {dparams['clip_percentile']} percentile for clipping before discretization.")
+
             self.discretize(clip_percentile=dparams['clip_percentile'], 
                             num_init=dparams['num_init'])
 
     def program_regulator_discovery(self, 
-                                    program_representation_raw: np.ndarray, 
+                                    program_representation: np.ndarray, 
                                     sample_id_key: str = 'sample_id',
                                     params: dict = None):
         """
@@ -414,7 +518,7 @@ class GeneDSPIN(AbstractDSPIN):
 
         Parameters
         ----------
-        program_representation_raw : np.ndarray
+        program_representation : np.ndarray
             The gene program representation.
         sample_id_key : str, optional
             Column name in `adata.obs` specifying sample identifiers. Default is 'sample_id'.
@@ -422,7 +526,7 @@ class GeneDSPIN(AbstractDSPIN):
             Additional parameters for regression. Default is None.
         """
 
-        program_states, samp_list = sample_states(self.adata.obs[sample_id_key], program_representation_raw)
+        program_states, samp_list = sample_states(self.adata.obs[sample_id_key], program_representation)
 
         gene_states, _ = sample_states(self.adata.obs[sample_id_key], self._onmf_rep_tri)
 
@@ -430,7 +534,9 @@ class GeneDSPIN(AbstractDSPIN):
             'num_epoch': 500,
             'stepsz': 0.01,
             'lambda_l1_interaction': 0.01,
-            'rec_gap': 10
+            'rec_gap': 10,
+            'backtrack_gap': 40,
+            'backtrack_tol': 5
         }
 
         if params is not None:
@@ -609,9 +715,9 @@ class ProgramDSPIN(AbstractDSPIN):
             Parameters for discretization of the oNMF representation. Default is an empty dictionary.
         """
 
-        if mode not in ['compute_summary', 'compute_only', 'summary_only']:
+        if mode not in ['compute_summary', 'compute_only', 'summary_only', 'precomputed']:
             raise ValueError(
-                'Mode must be one of compute_summary, compute_only, or summary_only')
+                'Mode must be one of compute_summary, compute_only, summary_only, or precomputed.')
         
         adata = self.adata
         num_spin = self.num_spin
@@ -631,7 +737,7 @@ class ProgramDSPIN(AbstractDSPIN):
         
         if onmf_parameters['num_subsample_large'] is None:
             onmf_parameters['num_subsample_large'] = min(onmf_parameters['num_subsample'] * 5, adata.shape[0])
-            print(onmf_parameters['num_subsample_large'])
+            # print(onmf_parameters['num_subsample_large'])
 
         if onmf_parameters['balance_method'] not in ['equal', 'proportional', 'squareroot', None]:
             raise ValueError('balance_method must be one of equal, proportional, squareroot, or None')
@@ -664,11 +770,11 @@ class ProgramDSPIN(AbstractDSPIN):
         self.prior_programs_ind = prior_program_ind
 
         num_onmf_components = onmf_parameters['num_onmf_components']
+
+        os.makedirs(self.save_path + 'onmf/', exist_ok=True)
         
         if mode == 'compute_summary' or mode == 'compute_only':
 
-            os.makedirs(self.save_path + 'onmf/', exist_ok=True)
-
             cur_gene_matrix = adata.X[:, prior_program_mask]
             cluster_label_raw = adata.obs[cluster_key].values
             cluster_label = pd.factorize(cluster_label_raw)[0]
@@ -681,7 +787,7 @@ class ProgramDSPIN(AbstractDSPIN):
 
                 compute_onmf_decomposition(sub_gene_matrix_normed, num_onmf_components, seed=cur_seed, params=onmf_parameters, save_path=self.save_path + f'onmf/onmf_components_{num_onmf_components}_repeat_{cur_seed}.npy')
             
-        if mode == 'compute_summary' or mode == 'summary_only':
+        if mode == 'compute_summary' or mode == 'summary_only' or mode == 'precomputed':
             
             cluster_label_raw = adata.obs[cluster_key].values
             cluster_label = pd.factorize(cluster_label_raw)[0]
@@ -689,20 +795,31 @@ class ProgramDSPIN(AbstractDSPIN):
             # Perform subsampling and standard deviation clipping on the matrix
             std_clipped, sub_large_gene_matrix_normed = subsample_normalize_gene_matrix(adata.X, cluster_label, num_subsample=onmf_parameters['num_subsample_large'], seed=seed, params=onmf_parameters)
             self.matrix_std = std_clipped
-            
-            # Summarize the oNMF decompositions
-            if num_onmf_components > 0:
-                onmf_summary = self.summarize_onmf_result(sub_large_gene_matrix_normed, num_onmf_components, num_repeat)
+
+            if mode != 'precomputed':
+                # Summarize the oNMF decompositions
+                if num_onmf_components > 0:
+                    onmf_summary = self.summarize_onmf_result(sub_large_gene_matrix_normed, num_onmf_components, num_repeat)
+                else:
+                    onmf_summary = self.compute_onmf_expression(sub_large_gene_matrix_normed, prior_program_ind)
+                
+                onmf_save_path = self.save_path + f'onmf/onmf_summary_total_{num_spin}_onmf_{num_onmf_components}_prior_{len(prior_program_ind)}.npy'
+                np.save(onmf_save_path, onmf_summary)
+                print(f"Consensus oNMF results saved to {onmf_save_path}")
             else:
-                onmf_summary = self.compute_onmf_expression(sub_large_gene_matrix_normed, prior_program_ind)
+                assert 'precomputed_onmf_path' in params.keys(), "For precomputed mode, 'precomputed_onmf_path' must be provided in params."
+                onmf_summary = np.load(params['precomputed_onmf_path'], allow_pickle=True).item()
             
             self._onmf_summary = onmf_summary
 
             # Save the gene programs to a CSV file
             file_path = onmf_to_csv(onmf_summary.components_, adata.var_names, self.save_path + f'gene_programs_total_{num_spin}_onmf_{num_onmf_components}_prior_{len(prior_program_ind)}.csv', thres=0.01)
-            print('Gene programs saved to {}'.format(file_path))
+            print('Gene program information saved to {}'.format(file_path))
             self.gene_program_csv = file_path
 
+            file_wpath = onmf_to_csv(onmf_summary.components_, adata.var_names, self.save_path + f'gene_programs_total_{num_spin}_onmf_{num_onmf_components}_prior_{len(prior_program_ind)}_weights.csv', thres=0.01, if_write_weights=True)
+            # print('Gene programs saved to {}'.format(file_path))
+            
             gene_matrix = self.adata.X
             if issparse(gene_matrix):
                 gene_matrix = np.asarray(gene_matrix.toarray()).astype(np.float64)
@@ -717,6 +834,20 @@ class ProgramDSPIN(AbstractDSPIN):
             self.discretize(clip_percentile=dparams['clip_percentile'],
                             num_init=dparams['num_init'])
 
+        num_display_gene = 5
+        gene_pd = pd.read_csv(self.gene_program_csv, index_col=None)
+        
+        top_gene_list = [gene_pd.iloc[:, ii].tolist()[:num_display_gene] for ii in range(num_spin)]
+        # remove nan
+        top_gene_list_filtered = []
+        for sub_list in top_gene_list:
+            sub_list_filtered = [gene for gene in sub_list if str(gene) != 'nan']
+            top_gene_list_filtered.append(sub_list_filtered)
+
+        program_names = [f'P{ii}-' + ','.join(top_gene_list_filtered[ii]) for ii in range(num_spin)]
+        self.name_list = program_names
+        self.name_list_short = [name.split('-')[0] for name in program_names]
+
 
 class DSPIN(object):
     """

{dspin-1.3.4 → dspin-1.4.0}/dspin/plot.py

@@ -9,19 +9,18 @@ import matplotlib.patheffects as patheffects
 from sklearn.cluster import KMeans
 import numpy as np
 import pandas as pd
-import seaborn as sns
-import os
-import sys
 import csv
 import scanpy as sc
 import igraph as ig
 import leidenalg as la
 import networkx as nx
 from scipy import optimize
-from scipy.spatial.distance import pdist
+from scipy.spatial.distance import pdist, squareform
+import matplotlib as mpl
+from typing import List, Dict, Optional, Tuple
 
 
-def onmf_to_csv(features, gene_name, file_path, thres=0.01):
+def onmf_to_csv(features, gene_name, file_path, thres=0.01, if_write_weights=False):
     """
     Extract gene names in ONMF to CSV.
 
@@ -30,6 +29,7 @@ def onmf_to_csv(features, gene_name, file_path, thres=0.01):
     gene_name (List[str]): Names of the genes corresponding to feature indices.
     file_path (str): Path of the file where the CSV will be written.
     thres (float): Threshold value for filtering features.
+    if_write_weights (bool): If True, writes weights to the CSV file.
 
     Returns:
     str: Path of the created CSV file.
@@ -51,6 +51,11 @@ def onmf_to_csv(features, gene_name, file_path, thres=0.01):
             cur_line.insert(0, spin)
             write.writerow(cur_line)  # Write the current line to the CSV file
 
+            if if_write_weights:
+                cur_line = [features[spin, ind].round(6) for ind in gene_ind]
+                cur_line.insert(0, f'Weights_{spin}')
+                write.writerow(cur_line)
+
     # Read the created CSV, transpose it and save it again
     pd_csv = pd.read_csv(file_path, names=range(rec_max_gene + 1))
     pd_csv = pd_csv.transpose()
@@ -551,3 +556,259 @@ def plot_final(
     ax.set_title(title, fontsize=node_fontsize * 1.5, y = 1.05)
 
 
+def plot_module_heatmap(j_mat: np.ndarray, 
+                       modules: List[List[int]], 
+                       spin_name_list: Optional[List[str]] = None,
+                       fig_folder: Optional[str] = None) -> None:
+    """
+    Plot a heatmap of the network adjacency matrix with modules highlighted.
+    
+    Parameters
+    ----------
+    j_mat : np.ndarray
+        The network adjacency matrix.
+    modules : List[List[int]]
+        List of modules, where each module is a list of node indices.
+    spin_name_list : List[str], optional
+        List of names for the spins/nodes. If None, will use generic names.
+    fig_folder : str, optional
+        Folder path to save the figure. If None, will not save the figure.
+    """
+    num_spin = j_mat.shape[0]
+    j_mat = j_mat.copy()
+    np.fill_diagonal(j_mat, 0)
+    
+    # Create spin order based on modules
+    spin_order = [spin for cur_list in modules for spin in cur_list]
+    net_class_len = [len(cur_list) for cur_list in modules]
+    
+    # Create custom colormap
+    cmap_hvec = mpl.colors.LinearSegmentedColormap.from_list("", ['#E84B23', '#FFFFFF', '#3285CC'])
+    
+    # Plot heatmap
+    sc.set_figure_params(figsize=[0.22 * num_spin + 0.5, 0.22 * num_spin])
+    plt.imshow(j_mat[:, spin_order][spin_order, :], cmap=cmap_hvec, vmin=-1, vmax=1)
+    
+    # Add module boundaries
+    for ii in range(len(net_class_len)):
+        plt.axhline(np.sum(net_class_len[:ii]) - 0.5, color='k', linewidth=1)
+        plt.axvline(np.sum(net_class_len[:ii]) - 0.5, color='k', linewidth=1)
+    
+    # Set labels
+    if spin_name_list is not None:
+        plt.xticks(np.arange(len(spin_order)), np.array(spin_name_list)[spin_order], rotation=90, fontsize=10)
+        plt.yticks(np.arange(len(spin_order)), np.array(spin_name_list)[spin_order], fontsize=10)
+    else:
+        plt.xticks(np.arange(len(spin_order)), [f'P{spin}' for spin in spin_order], rotation=90, fontsize=10)
+        plt.yticks(np.arange(len(spin_order)), [f'P{spin}' for spin in spin_order], fontsize=10)
+        
+    plt.colorbar(fraction=0.03)
+    plt.grid()
+    plt.title(f'Network on {num_spin} programs', fontsize=12)
+    
+    # Save figure if folder is provided
+    if fig_folder is not None:
+        plt.savefig(fig_folder + 'module_heatmap.png', dpi=300, bbox_inches='tight')
+    
+    plt.show()
+
+
+
+def layout_loss_fun(xx, ori_pos, network):
+
+    stength_module = 12
+    strength_repulsion = 12
+    strength_spring = 1
+
+    spring_length = 2
+    spring_weight = np.abs(network)
+
+    module_length = 1
+
+    xx = xx.reshape(- 1, 2)
+    module_dist = np.clip(np.sqrt(((xx - ori_pos) ** 2).sum(axis=1)) - module_length, 0, None)
+    module_attract = np.sum(module_dist ** 2)
+    repulse = np.sum(1 / pdist(xx))
+
+    pairwise_distance = squareform(pdist(xx))
+    clip_dist = np.clip(pairwise_distance - spring_length, 0, None)
+    spring_energy = np.sum(clip_dist ** 2 * spring_weight)
+
+    return stength_module * module_attract + strength_repulsion * repulse + strength_spring * spring_energy
+
+def compute_module_angle(module_size, gap_size):
+
+    module_plots_size = np.sqrt(module_size)
+
+    angle_list = []
+    cur_angle = 0
+
+    for ii in range(len(module_plots_size)):
+        angle_list.append(cur_angle)
+        cur_angle += module_plots_size[ii] + gap_size
+
+    angle_list = np.array(angle_list)
+    angle_list /= (cur_angle / (2 * np.pi))
+
+    return angle_list
+
+def compute_pos_from_module(network, net_class, circle_size=None, module_angle=None, gap_size=0.5):
+
+    module_size = np.array([len(cur_list) for cur_list in net_class])
+    num_spin = np.sum(module_size)
+
+    if module_angle is None:
+        module_angle = compute_module_angle(module_size, gap_size)
+
+    if circle_size is None:
+        circle_size = num_spin ** 0.3
+
+    module_center = circle_size * np.array([np.cos(module_angle), np.sin(module_angle)]).T
+
+    pos_module = np.zeros((num_spin, 2))
+
+    for ii in range(len(module_size)):
+        cur_list = net_class[ii]
+        pos_module[cur_list, :] = module_center[ii, :]
+
+    pos_module += np.random.normal(0, 0.05, pos_module.shape)
+
+    ori_pos = pos_module.copy()
+
+    opt_res = optimize.minimize(lambda xx : layout_loss_fun(xx, ori_pos, network), pos_module.flatten())
+    pos_raw = opt_res.x.reshape(- 1, 2)
+
+    pos = {ii: pos_raw[ii] for ii in range(num_spin)}
+
+    return pos
+
+
+def compute_network_positions(j_mat: np.ndarray,
+                             modules: List[List[int]])-> Dict[int, Tuple[float, float]]:
+    """
+    Compute positions for network nodes based on modules.
+    
+    Parameters
+    ----------
+    j_mat : np.ndarray
+        The network adjacency matrix.
+    modules : List[List[int]]
+        List of modules, where each module is a list of node indices.
+    
+    Returns
+    -------
+    Dict[int, Tuple[float, float]]
+        Dictionary mapping node indices to (x, y) positions.
+    """
+    
+    pos = compute_pos_from_module(j_mat, modules)
+    
+    return pos
+
+
+def plot_network_diagram(j_mat: np.ndarray,
+                        modules: List[List[int]],
+                        pos: Optional[Dict[int, Tuple[float, float]]] = None, 
+                        spin_name_list_short: Optional[List[str]] = None,
+                        spin_name_list: Optional[List[str]] = None,
+                        fig_folder: Optional[str] = None,
+                        num_spin: Optional[int] = None) -> None:
+    """
+    Plot the network diagram with modules highlighted.
+    
+    Parameters
+    ----------
+    j_mat : np.ndarray
+        The network adjacency matrix.
+    modules : List[List[int]]
+        List of modules, where each module is a list of node indices.
+    pos : Dict[int, Tuple[float, float]], optional
+        Dictionary mapping node indices to (x, y) positions. If None, positions will be computed automatically.
+    spin_name_list_short : List[str], optional
+        Short names for the spins/nodes (used for labels). If None, will use generic names.
+    spin_name_list : List[str], optional
+        Full names for the spins/nodes (used for tooltips). If None, will use generic names.
+    fig_folder : str, optional
+        Folder path to save the figure. If None, will not save the figure.
+    num_spin : int, optional
+        Number of spins. If None, will be inferred from j_mat shape.
+    """
+    if num_spin is None:
+        num_spin = j_mat.shape[0]
+    
+    # Filter network for visualization
+    thres = 0.2
+    j_filt = j_mat.copy()
+    j_filt[np.abs(j_mat) < thres] = 0
+    np.fill_diagonal(j_filt, 0)
+    
+    # Create networkx graph
+    G = nx.from_numpy_array(j_filt)
+    
+    # Compute positions
+    if pos is None:
+        pos = compute_network_positions(j_mat, modules)
+    
+    # Create labels
+    if spin_name_list_short is not None:
+        label = {ii: spin_name_list_short[ii] for ii in range(num_spin)}
+    else:
+        label = {ii: f'P{ii}' for ii in range(num_spin)}
+    
+    # Plot network
+    fig_size = 3 * num_spin ** 0.3
+    sc.set_figure_params(figsize=[fig_size, fig_size])
+    fig, ax = plt.subplots()
+    
+    # Draw edges
+    eposi = [(u, v) for (u, v, d) in G.edges(data=True) if d['weight'] > 0]
+    wposi = np.array([d['weight'] for (u, v, d) in G.edges(data=True) if d['weight'] > 0])
+    
+    enega = [(u, v) for (u, v, d) in G.edges(data=True) if d['weight'] < 0]
+    wnega = np.array([d['weight'] for (u, v, d) in G.edges(data=True) if d['weight'] < 0])
+    
+    # Draw positive edges
+    if len(eposi) > 0:
+        nx.draw_networkx_edges(G, pos, ax=ax, edgelist=eposi, width=3 * wposi,
+                               edge_color='#3285CC', alpha=0.7)
+    
+    # Draw negative edges
+    if len(enega) > 0:
+        nx.draw_networkx_edges(G, pos, ax=ax, edgelist=enega, width=3 * (-wnega),
+                               edge_color='#E84B23', alpha=0.7)
+    
+    # Draw nodes
+    node_colors = ['#f0dab1'] * num_spin
+    nx.draw_networkx_nodes(G, pos, ax=ax, node_size=300, node_color=node_colors, 
+                          edgecolors='black', linewidths=1)
+    
+    # Draw labels
+    for spin in range(num_spin):
+        plt.text(pos[spin][0], pos[spin][1], label[spin], fontsize=10, 
+                path_effects=[patheffects.withStroke(linewidth=3, foreground='w')], 
+                ha='center', va='center')
+    
+    ax.set_aspect('equal')
+    ax.set_axis_off()
+    plt.title('Network Diagram with Modules', fontsize=12)
+
+    # Print the spin_name_list in a column to the right of the figure
+    if spin_name_list is not None:
+        # Place the text to the right of the plot
+        # Get axis limits to determine placement
+        xlim = ax.get_xlim()
+        ylim = ax.get_ylim()
+        # Place the column at a little to the right of the rightmost node
+        x_text = xlim[1] + 0.02 * (xlim[1] - xlim[0])
+        y_start = ylim[1] - 0.05 * (ylim[1] - ylim[0])
+        line_gap = fig_size / 20
+        for idx, name in enumerate(spin_name_list):
+            ax.text(x_text, y_start - idx * line_gap, name, fontsize=10, va='top', ha='left')
+        # Optionally, expand the xlim to make space for the text
+        ax.set_xlim(xlim[0], x_text + 1.5)
+    
+    # Save figure if folder is provided
+    if fig_folder is not None:
+        plt.savefig(fig_folder + 'network_diagram.png', dpi=300, bbox_inches='tight')
+    
+    plt.show() 

{dspin-1.3.4 → dspin-1.4.0/dspin.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: dspin
-Version: 1.3.4
+Version: 1.4.0
 Summary: Regulatory network models from single-cell perturbation profiling
 Home-page: https://github.com/JialongJiang/DSPIN
 Author: Jialong Jiang
@@ -25,6 +25,18 @@ Requires-Dist: scanpy
 Requires-Dist: tqdm
 Requires-Dist: igraph
 Requires-Dist: leidenalg
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
 
 # D-SPIN: Constructing regulatory network models from scRNA-seq perturbation profiling  
 

{dspin-1.3.4 → dspin-1.4.0}/setup.py

@@ -21,7 +21,7 @@ install_requires = [
 # ]
 
 # Managing your version (adjust accordingly)
-version = "1.3.4"  # You can manage versioning in a more sophisticated way if needed
+version = "1.4.0"  # You can manage versioning in a more sophisticated way if needed
 
 # Long description (usually a README file)
 readme = open("README_PyPI.md").read()  # Adjust the file name if necessary

dspin-1.3.4/dspin/__init__.py

@@ -1,5 +0,0 @@
-
-from . import annotate
-from . import compute
-from . import plot
-from . import dspin