Merge pull request #2228 from broadinstitute/jb-de-cluster-name-ref

Make updating existing DE result more robust (SCP-5952)

Author: bistline

app/lib/differential_expression_service.rb

@@ -134,6 +134,7 @@ def self.run_differential_expression_job(cluster_group, study, user, annotation_
     metadata_url = study_file.is_viz_anndata? ?
                      RequestUtils.data_fragment_url(study_file, 'metadata') :
                      study.metadata_file.gs_url
+    cluster_group_id = cluster_group.id
     # begin assembling parameters
     de_params = {
       annotation_name:,
@@ -143,10 +144,12 @@ def self.run_differential_expression_job(cluster_group, study, user, annotation_
       group2:,
       annotation_file: annotation_scope == 'cluster' ? cluster_url : metadata_url,
       cluster_file: cluster_url,
-      cluster_name: cluster_group.name
+      cluster_name: set_cluster_name(study, cluster_group, annotation_name, annotation_scope),
+      cluster_group_id:
     }
     raw_matrix = ClusterVizService.raw_matrix_for_cluster_cells(study, cluster_group)
     de_params[:matrix_file_path] = raw_matrix.gs_url
+    de_params[:matrix_file_id] = raw_matrix.id
     if raw_matrix.file_type == 'MM Coordinate Matrix'
       de_params[:matrix_file_type] = 'mtx'
       # we know bundle exists and is completed as :raw_matrix_for_cluster_cells will throw an exception if it isn't
@@ -287,6 +290,21 @@ def self.annotation_eligible?(name)
     ALLOWED_ANNOTS =~ name && EXCLUDED_ANNOTS !~ name
   end
 
+  # determine if the requested cluster/annotation has an existing result object
+  # used for setting cluster name or in pairwise DE when merging in new results
+  #
+  # * *params*
+  #   - +study+            (Study) => Study in which results exist
+  #   - +cluster_group+    (ClusterGroup) => Clustering object to source name/file from
+  #   - +annotation_name+  (String) => Name of requested annotation
+  #   - +annotation_scope+ (String) => Scope of requested annotation ('study' or 'cluster')
+  #
+  # * *returns*
+  #   - (DifferentialExpressionResult, nil)
+  def self.find_existing_result(study, cluster_group, annotation_name, annotation_scope)
+    DifferentialExpressionResult.find_by(study:, cluster_group:, annotation_name:, annotation_scope:)
+  end
+
   # determine if a study already has DE results for an annotation, taking scope into account
   # cluster-based annotations must match to the specified cluster in the annotation object
   # for study-wide annotations, return true if any results exist, regardless of cluster as this indicates that DE
@@ -309,6 +327,21 @@ def self.results_exist?(study, annotation)
     ).exists?
   end
 
+  # helper to set cluster_name when existing results are present
+  # this prevents result file URLs breaking because the cluster has been renamed at some point
+  #
+  # * *params*
+  #   - +study+            (Study) => Study in which results exist
+  #   - +cluster_group+    (ClusterGroup) => Clustering object to source name/file from
+  #   - +annotation_name+  (String) => Name of requested annotation
+  #   - +annotation_scope+ (String) => Scope of requested annotation ('study' or 'cluster')
+  #
+  # * *returns*
+  #   - (String)
+  def self.set_cluster_name(study, cluster_group, annotation_name, annotation_scope)
+    find_existing_result(study, cluster_group, annotation_name, annotation_scope)&.cluster_name || cluster_group.name
+  end
+
   # determine if a study meets the requirements for differential expression:
   # 1. public
   # 2. has clustering/metadata
@@ -383,12 +416,15 @@ def self.validate_annotation(cluster_group, study, annotation_name, annotation_s
     elsif pairwise
       missing = [group1, group2] - annotation[:values]
       raise ArgumentError, "#{annotation_name} does not contain '#{missing.join(', ')}'" if missing.any?
+
       cell_count = {
-        "#{group1}" => cells_by_label[group1].count,
-        "#{group2}" => cells_by_label[group2].count
+        group1.to_s => cells_by_label[group1]&.count.to_i,
+        group2.to_s => cells_by_label[group2]&.count.to_i
       }.keep_if { |_, c| c < 2 }
-      raise ArgumentError,
-            "#{cell_count.keys.join(', ')} does not have enough cells represented in #{identifier}" if cell_count.any?
+      if cell_count.any?
+        raise ArgumentError, "#{cell_count.keys.join(', ')} does not have enough cells represented in #{identifier} " \
+          "for #{cluster_group.name}"
+      end
     end
   end
 

app/models/differential_expression_parameters.rb

@@ -17,12 +17,14 @@ class DifferentialExpressionParameters
   # annotation_scope: scope of annotation (study, cluster)
   # annotation_file: source file for above annotation
   # cluster_file: clustering file with cells to use as control list for DE
-  # cluster_name: name of associated ClusterGroup object
+  # cluster_name: name of associated ClusterGroup object for use in output filenames
+  # cluster_group_id: BSON ID of ClusterGroup object for associations
   # de_type: type of differential expression calculation: 'rest' (one-vs-rest) or 'pairwise'
   # group1: first group for pairwise comparison
   # group2: second group for pairwise comparison
   # matrix_file_path: raw counts matrix with source expression data
   # matrix_file_type: type of raw counts matrix (dense, sparse)
+  # matrix_file_id: BSON ID of raw matrix StudyFile for associations
   # gene_file (optional): genes/features file for sparse matrix
   # barcode_file (optional): barcodes file for sparse matrix
   # machine_type (optional): override for default ingest machine type (uses 'n2d-highmem-8')
@@ -34,24 +36,26 @@ class DifferentialExpressionParameters
     annotation_file: nil,
     cluster_file: nil,
     cluster_name: nil,
+    cluster_group_id: nil,
     de_type: 'rest',
     group1: nil,
     group2: nil,
     matrix_file_path: nil,
     matrix_file_type: nil,
+    matrix_file_id: nil,
     gene_file: nil,
     barcode_file: nil,
     machine_type: 'n2d-highmem-8',
     file_size: 0
   }.freeze
 
   # values that are available as methods but not as attributes (and not passed to command line)
-  NON_ATTRIBUTE_PARAMS = %i[machine_type file_size].freeze
+  NON_ATTRIBUTE_PARAMS = %i[machine_type file_size cluster_group_id matrix_file_id].freeze
 
   attr_accessor(*PARAM_DEFAULTS.keys)
 
-  validates :annotation_name, :annotation_scope, :annotation_file, :cluster_file,
-            :cluster_name, :matrix_file_path, :matrix_file_type, presence: true
+  validates :annotation_name, :annotation_scope, :annotation_file, :cluster_file, :cluster_name, :cluster_group_id,
+            :matrix_file_path, :matrix_file_type, :matrix_file_id, presence: true
   validates :annotation_file, :cluster_file, :matrix_file_path,
             format: { with: Parameterizable::GS_URL_REGEXP, message: 'is not a valid GS url' }
   validates :annotation_scope, inclusion: %w[cluster study]
@@ -75,4 +79,12 @@ def initialize(attributes = nil)
       self.machine_type = assign_machine_type
     end
   end
+
+  def cluster_group
+    ClusterGroup.find(cluster_group_id)
+  end
+
+  def matrix_file
+    StudyFile.find(matrix_file_id)
+  end
 end

app/models/differential_expression_result.rb

@@ -203,7 +203,7 @@ def set_automated_comparisons(group1: nil, group2: nil)
   private
 
   def set_cluster_name
-    self.cluster_name = cluster_group.name
+    self.cluster_name = cluster_group.name unless cluster_name
   end
 
   def comparisons_available?

app/models/ingest_job.rb

@@ -738,14 +738,13 @@ def launch_differential_expression_jobs
   # set corresponding differential expression flags on associated annotation
   def create_differential_expression_results
     annotation_identifier = "#{params_object.annotation_name}--group--#{params_object.annotation_scope}"
-    Rails.logger.info "Creating differential expression result object for annotation: #{annotation_identifier}"
-    cluster = ClusterGroup.find_by(study_id: study.id, study_file_id: study_file.id, name: params_object.cluster_name)
-    matrix_url = params_object.matrix_file_path
-    bucket_path = matrix_url.split("gs://#{study.bucket_id}/").last
-    matrix_file = StudyFile.where(
-      study:, :upload_file_name.in => [bucket_path, bucket_path.split('/').last]
-    ).detect(&:is_raw_counts_file?)
-    de_result = DifferentialExpressionResult.find_or_initialize_by(
+    cluster = params_object.cluster_group
+    matrix_file = params_object.matrix_file
+    Rails.logger.info "Creating differential expression result object for #{annotation_identifier} " \
+                        "(cluster: #{cluster.name} in #{study.accession})"
+    de_result = DifferentialExpressionService.find_existing_result(
+      study, cluster, params_object.annotation_name, params_object.annotation_scope
+    ) || DifferentialExpressionResult.new(
       study: study, cluster_group: cluster, cluster_name: cluster.name,
       annotation_name: params_object.annotation_name, annotation_scope: params_object.annotation_scope,
       matrix_file_id: matrix_file.id
@@ -1000,7 +999,7 @@ def get_job_analytics
         )
       end
     when :differential_expression
-      cluster = ClusterGroup.find_by(study_id: study.id, study_file_id: study_file.id, name: params_object.cluster_name)
+      cluster = params_object.cluster_group
       annotation_params = {
         cluster: cluster,
         annot_name: params_object.annotation_name,

test/models/batch_api_client_test.rb

@@ -117,14 +117,17 @@ class BatchApiClientTest < ActiveSupport::TestCase
   test 'should create and submit DE Batch API job' do
     action = :differential_expression
     bucket_dir = "_scp_internal/anndata_ingest/#{@ann_data_study.accession}_#{@ann_data_file.id}"
+    cluster = @ann_data_study.cluster_groups.first
     anndata_options = {
       annotation_name: 'disease',
       annotation_scope: 'study',
       annotation_file: "gs://#{@ann_data_study.bucket_id}/#{bucket_dir}/h5ad_frag.metadata.tsv.gz",
       cluster_file: "gs://#{@ann_data_study.bucket_id}/#{bucket_dir}/h5ad_frag.cluster.X_umap.tsv.gz",
       cluster_name: 'umap',
+      cluster_group_id: cluster.id,
       matrix_file_path: @ann_data_file.gs_url,
       matrix_file_type: 'h5ad',
+      matrix_file_id: @ann_data_file.id,
       file_size: @ann_data_file.upload_file_size
     }
     params_object = DifferentialExpressionParameters.new(**anndata_options)

test/models/differential_expression_parameters_test.rb

@@ -3,14 +3,18 @@
 class DifferentialExpressionParametersTest < ActiveSupport::TestCase
 
   before(:all) do
+    cluster_group_id = BSON::ObjectId.new
+    matrix_file_id = BSON::ObjectId.new
     @dense_options = {
       annotation_name: 'Category',
       annotation_scope: 'cluster',
       annotation_file: 'gs://test_bucket/metadata.tsv',
       cluster_file: 'gs://test_bucket/cluster.tsv',
       cluster_name: 'UMAP',
+      cluster_group_id:,
       matrix_file_path: 'gs://test_bucket/dense.tsv',
-      matrix_file_type: 'dense'
+      matrix_file_type: 'dense',
+      matrix_file_id:
     }
 
     @sparse_options = {
@@ -19,8 +23,10 @@ class DifferentialExpressionParametersTest < ActiveSupport::TestCase
       annotation_file: 'gs://test_bucket/metadata.tsv',
       cluster_file: 'gs://test_bucket/cluster.tsv',
       cluster_name: 'UMAP',
+      cluster_group_id:,
       matrix_file_path: 'gs://test_bucket/sparse.tsv',
       matrix_file_type: 'mtx',
+      matrix_file_id:,
       gene_file: 'gs://test_bucket/genes.tsv',
       barcode_file: 'gs://test_bucket/barcodes.tsv'
     }
@@ -31,8 +37,10 @@ class DifferentialExpressionParametersTest < ActiveSupport::TestCase
       annotation_file: 'gs://test_bucket/metadata.tsv',
       cluster_file: 'gs://test_bucket/cluster.tsv',
       cluster_name: 'UMAP',
+      cluster_group_id:,
       matrix_file_path: 'gs://test_bucket/matrix.h5ad',
       matrix_file_type: 'h5ad',
+      matrix_file_id:,
       file_size: 10.gigabytes
     }
 
@@ -45,8 +53,10 @@ class DifferentialExpressionParametersTest < ActiveSupport::TestCase
       annotation_file: 'gs://test_bucket/metadata.tsv',
       cluster_file: 'gs://test_bucket/cluster.tsv',
       cluster_name: 'UMAP',
+      cluster_group_id:,
       matrix_file_path: 'gs://test_bucket/dense.tsv',
-      matrix_file_type: 'dense'
+      matrix_file_type: 'dense',
+      matrix_file_id:
     }
   end
 
@@ -144,4 +154,47 @@ class DifferentialExpressionParametersTest < ActiveSupport::TestCase
     dense_params = DifferentialExpressionParameters.new(@dense_options)
     assert_not_includes dense_params.attributes, :machine_type
   end
+
+  test 'should find cluster group and matrix file' do
+    user = FactoryBot.create(:user, test_array: @@users_to_clean)
+    study = FactoryBot.create(:detached_study,
+                              name_prefix: 'DifferentialExpressionParameters Test',
+                              user:,
+                              test_array: @@studies_to_clean)
+    raw_matrix = FactoryBot.create(:expression_file,
+                                   name: 'raw.txt',
+                                   study:,
+                                   expression_file_info: {
+                                     is_raw_counts: true,
+                                     units: 'raw counts',
+                                     library_preparation_protocol: 'Drop-seq',
+                                     biosample_input_type: 'Whole cell',
+                                     modality: 'Proteomic'
+                                   })
+
+    cells = %w[A B C D E]
+    cluster_file = FactoryBot.create(:cluster_file,
+                                     name: 'cluster_diffexp.tsv',
+                                     study:,
+                                     cell_input: {
+                                       x: [1, 4, 6, 8, 9],
+                                       y: [7, 5, 3, 2, 1],
+                                       cells:
+                                     })
+    cluster_group = ClusterGroup.find_by(study:, study_file: cluster_file)
+    params = {
+      annotation_name: 'Category',
+      annotation_scope: 'cluster',
+      annotation_file: 'gs://test_bucket/metadata.tsv',
+      cluster_file: 'gs://test_bucket/cluster_diffexp.tsv',
+      cluster_name: 'cluster_diffexp.tsv',
+      cluster_group_id: cluster_group.id,
+      matrix_file_path: 'gs://test_bucket/dense.tsv',
+      matrix_file_type: 'dense',
+      matrix_file_id: raw_matrix.id
+    }
+    de_params = DifferentialExpressionParameters.new(params)
+    assert_equal cluster_group, de_params.cluster_group
+    assert_equal raw_matrix, de_params.matrix_file
+  end
 end

test/models/ingest_job_test.rb

@@ -879,7 +879,8 @@ class IngestJobTest < ActiveSupport::TestCase
     cluster_file = "gs://#{bucket}/anndata/h5ad_frag.cluster.X_umap.tsv"
     params_object = DifferentialExpressionParameters.new(
       matrix_file_path: "gs://#{bucket}/matrix.h5ad", matrix_file_type: 'h5ad', file_size: study_file.upload_file_size,
-      annotation_file:, cluster_file:, cluster_name: 'umap', annotation_name: 'disease', annotation_scope: 'study'
+      annotation_file:, cluster_file:, cluster_name: 'umap', annotation_name: 'disease', annotation_scope: 'study',
+      cluster_group_id: BSON::ObjectId.new, matrix_file_id: study_file.id
     )
     pipeline_name = SecureRandom.uuid
     study_file.update(parse_status: 'parsing')
@@ -1079,7 +1080,7 @@ class IngestJobTest < ActiveSupport::TestCase
     # one vs rest test
     one_vs_rest = DifferentialExpressionParameters.new(
       annotation_name: 'cell_type__ontology_label', annotation_scope: 'study', cluster_name: 'cluster_diffexp.txt',
-      matrix_file_path: "gs://#{study.bucket_id}/raw.txt"
+      matrix_file_path: "gs://#{study.bucket_id}/raw.txt", cluster_group_id: cluster_group.id, matrix_file_id: matrix.id
     )
     job = IngestJob.new(study:, study_file: cluster_file, action: :differential_expression, params_object: one_vs_rest)
     job.create_differential_expression_results
@@ -1094,7 +1095,8 @@ class IngestJobTest < ActiveSupport::TestCase
     # pairwise test
     pairwise = DifferentialExpressionParameters.new(
       annotation_name: 'cell_type__ontology_label', annotation_scope: 'study', cluster_name: 'cluster_diffexp.txt',
-      matrix_file_path: "gs://#{study.bucket_id}/raw.txt", de_type: 'pairwise', group1: 'B cell', group2: 'T cell'
+      matrix_file_path: "gs://#{study.bucket_id}/raw.txt", de_type: 'pairwise', group1: 'B cell', group2: 'T cell',
+      cluster_group_id: cluster_group.id, matrix_file_id: matrix.id
     )
     job = IngestJob.new(study:, study_file: cluster_file, action: :differential_expression, params_object: pairwise)
     job.create_differential_expression_results

test/services/differential_expression_service_test.rb

@@ -445,4 +445,19 @@ class DifferentialExpressionServiceTest < ActiveSupport::TestCase
     user.reload
     assert_equal 0, user.weekly_de_quota
   end
+
+  test 'should set cluster name' do
+    annotation_name = 'cell_type__ontology_label'
+    annotation_scope = 'study'
+    assert_equal @cluster_group.name, DifferentialExpressionService.set_cluster_name(
+      @basic_study, @cluster_group, annotation_name, annotation_scope
+    )
+    DifferentialExpressionResult.create(
+      study: @basic_study, cluster_group: @cluster_group, cluster_name: 'umap', annotation_name:, annotation_scope:,
+      matrix_file_id: @raw_matrix.id
+    )
+    assert_equal 'umap', DifferentialExpressionService.set_cluster_name(
+      @basic_study, @cluster_group, annotation_name, annotation_scope
+    )
+  end
 end