Use New Spatial Regularizer

nextflow.config

@@ -11,23 +11,22 @@ process {
     publishDir = { "${params.outdir}" }
 }
 
-dag {
+def trace_timestamp = new java.util.Date().format( 'yyyy-MM-dd_HH-mm-ss')
+timeline {
     enabled = true
-    file = "${params.outdir}/pipeline_info/pipeline_dag.html"
-    overwrite = true
+    file    = "${params.outdir}/pipeline_info/execution_timeline_${trace_timestamp}.html"
+}
+report {
+    enabled = true
+    file    = "${params.outdir}/pipeline_info/execution_report_${trace_timestamp}.html"
 }
-
 trace {
     enabled = true
-    file = "${params.outdir}/pipeline_info/pipeline_trace.txt"
-    fields = 'task_id,name,status,exit,realtime,%cpu,rss'
-    overwrite = true
+    file    = "${params.outdir}/pipeline_info/execution_trace_${trace_timestamp}.txt"
 }
-
-report {
+dag {
     enabled = true
-    file = "${params.outdir}/pipeline_info/pipeline_report.html"
-    overwrite = true
+    file    = "${params.outdir}/pipeline_info/pipeline_dag_${trace_timestamp}.html"
 }
 
 // Function to ensure that resource requirements don't go beyond

src/bayestme/svi/deconvolution.py

@@ -22,42 +22,6 @@
 logger = logging.getLogger(__name__)
 
 
-def construct_edge_adjacency(neighbors):
-    data, rows, cols = [], [], []
-    nrows = 0
-    for i, j in neighbors:
-        data.extend([1, -1])
-        rows.extend([nrows, nrows])
-        cols.extend([i, j])
-        nrows += 1
-    indices = torch.tensor([rows, cols])
-    values = torch.tensor(data)
-    edge_adjacency_matrix = torch.sparse_coo_tensor(indices, values).float()
-    return edge_adjacency_matrix
-
-
-def construct_trendfilter(adjacency_matrix, k):
-    transformed_edge_adjacency_matrix = adjacency_matrix.clone()
-    for i in range(k):
-        if i % 2 == 0:
-            transformed_edge_adjacency_matrix = adjacency_matrix.t().mm(
-                transformed_edge_adjacency_matrix
-            )
-        else:
-            transformed_edge_adjacency_matrix = adjacency_matrix.mm(
-                transformed_edge_adjacency_matrix
-            )
-    extra = torch.sparse_coo_tensor(
-        torch.tensor([[0], [0]]),
-        torch.tensor([1]),
-        size=(1, transformed_edge_adjacency_matrix.shape[1]),
-    )
-    transformed_edge_adjacency_matrix = torch.vstack(
-        [transformed_edge_adjacency_matrix, extra]
-    )
-    return transformed_edge_adjacency_matrix
-
-
 def rv_should_be_sampled(site):
     """
     Determine if a random variable in the model should be saved or
@@ -73,19 +37,6 @@ def rv_should_be_sampled(site):
     )
 
 
-def create_reads_trace(psi, exp_profile, exp_load, cell_num_total):
-    """
-    :param psi: <N tissue spots> x <N components> matrix
-    :param exp_profile: <N components> x <N markers> matrix
-    :param exp_load: <N components> matrix
-    :param cell_num_total: <N tissue spots> matrix
-    :return: <N tissue spots> x <N markers> x <N components> matrix
-    """
-    number_of_cells_per_component = (psi.T * cell_num_total).T * exp_load.T
-    result = number_of_cells_per_component[:, :, None] * exp_profile[None, :, :]
-    return np.transpose(result, (0, 2, 1))
-
-
 class BayesTME_VI:
     def __init__(
         self,
@@ -111,8 +62,6 @@ def __init__(
         self.N = stdata.n_spot_in
         self.n_genes = stdata.n_gene
         self.edges = get_edges(stdata.positions_tissue, layout=stdata.layout)
-        self.delta = construct_edge_adjacency(self.edges)
-        self.delta = construct_trendfilter(self.delta, 0)
         self.spatial_regularization_coefficient = rho
         self.losses = []
         self.expression_truth_weight = expression_truth_weight
@@ -246,7 +195,7 @@ def spatial_regularizer(self, x):
         # Delta should be of size (n_edges * n_celltype) by (n_spot * n_celltype)
         # x should be of size n_spot by n_celltype
         return (
-            torch.abs(self.Delta @ x.reshape(-1, 1)).sum()
+            torch.abs(x[self.edges[:, 0]] - x[self.edges[:, 1]]).sum()
             * self.spatial_regularization_coefficient
         )
 
@@ -258,8 +207,7 @@ def deconvolution(self, K, n_iter=10000, n_traces=1000, use_spatial_guide=True):
             )
         else:
             self.n_celltypes = K
-        # TODO: maybe make Delta sparse, but need to change spatial_regularizer as well (double check if sparse grad is supported)
-        self.Delta = torch.kron(torch.eye(self.n_celltypes), self.delta.to_dense())
+
         logger.info("Optimizer: {} {}".format("Adam", self.opt_params))
         logger.info(
             "Deconvolving: {} spots, {} genes, {} cell types".format(