update geneenrichment to be less strict (#1529)

### Description

<!-- Provide a detailed description of the changes in this PR -->

#### Usage

<!--- How does a user interact with the changed code -->

```python
TODO: Add code snippet
```

### Type of changes

<!-- Mark the relevant option with an [x] -->

- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Refactor
- [ ] Documentation update
- [ ] Other (please describe):

### CI Pipeline Configuration

Configure CI behavior by applying the relevant labels. By default, only
basic unit tests are run.

-
[ciflow:skip](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:skip)
- Skip all CI tests for this PR
-
[ciflow:notebooks](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:notebooks)
- Run Jupyter notebooks execution tests for bionemo2
-
[ciflow:slow](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:slow)
- Run slow single GPU integration tests marked as @pytest.mark.slow for
bionemo2
-
[ciflow:all](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:all)
- Run all tests (unit tests, slow tests, and notebooks) for bionemo2.
This label can be used to enforce running tests for all bionemo2.
-
[ciflow:all-recipes](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:all-recipes)
- Run tests for all recipes (under bionemo-recipes). This label can be
used to enforce running tests for all recipes.

Unit tests marked as `@pytest.mark.multi_gpu` or
`@pytest.mark.distributed` are not run in the PR pipeline.

For more details, see [CONTRIBUTING](CONTRIBUTING.md)

> [!NOTE]
> By default, only basic unit tests are run. Add appropriate labels to
enable an additional test coverage.

#### Authorizing CI Runs

We use
[copy-pr-bot](https://docs.gha-runners.nvidia.com/apps/copy-pr-bot/#automation)
to manage authorization of CI
runs on NVIDIA's compute resources.

- If a pull request is opened by a trusted user and contains only
trusted changes, the pull request's code will
automatically be copied to a pull-request/ prefixed branch in the source
repository (e.g. pull-request/123)
- If a pull request is opened by an untrusted user or contains untrusted
changes, an NVIDIA org member must leave an
`/ok to test` comment on the pull request to trigger CI. This will need
to be done for each new commit.

#### Triggering Code Rabbit AI Review

To trigger a code review from code rabbit, comment on a pull request
with one of these commands:

- @coderabbitai review - Triggers a standard review
- @coderabbitai full review - Triggers a comprehensive review

See https://docs.coderabbit.ai/reference/review-commands for a full list
of commands.

### Pre-submit Checklist

<!--- Ensure all items are completed before submitting -->

- [ ] I have tested these changes locally
- [ ] I have updated the documentation accordingly
- [ ] I have added/updated tests as needed
- [ ] All existing tests pass successfully

---------

Signed-off-by: jwilber <jwilber@nvidia.com>

Author: jwilber

bionemo-recipes/interpretability/sparse_autoencoders/recipes/codonfm/codon_dashboard/src/App.jsx

@@ -64,7 +64,7 @@ class FeatureTooltip {
   }
 }
 
-export default function EmbeddingView({ brush, categoryColumn, categoryColumns, onFeatureClick, highlightedFeatureId, viewportState, onViewportChange, labels, features, selectedCategory, darkMode }) {
+export default function EmbeddingView({ brush, categoryColumn, categoryColumns, onFeatureClick, highlightedFeatureId, viewportState, onViewportChange, labels, features, selectedCategory, darkMode, hiddenCategories }) {
   const containerRef = useRef(null)
   const viewRef = useRef(null)
   const onFeatureClickRef = useRef(onFeatureClick)
@@ -267,7 +267,8 @@ export default function EmbeddingView({ brush, categoryColumn, categoryColumns,
     if (!viewRef.current) return
 
     let categoryColName = null
-    let colors = Array(50).fill(DEFAULT_COLOR)
+    const HIDDEN_COLOR = darkMode ? "#0a0a0a" : "#fafafa"
+    let colors = Array(50).fill(HIDDEN_COLOR)
 
     if (categoryColumn && categoryColumn !== "none") {
       const colInfo = categoryColumns?.find(c => c.name === categoryColumn)
@@ -278,6 +279,17 @@ export default function EmbeddingView({ brush, categoryColumn, categoryColumns,
         } else if (colInfo.type === 'string') {
           categoryColName = `${categoryColumn}_cat`
           colors = CATEGORY_COLORS.slice(0, Math.max(colInfo.nUnique, 10))
+          // Map colors to match DENSE_RANK order, dim non-selected when filtering
+          if (hiddenCategories && hiddenCategories.size > 0 && features) {
+            const allCatNames = [...new Set(
+              features.map(f => f[categoryColumn]).filter(v => v != null)
+            )].sort()
+            colors = colors.map((c, i) => {
+              const name = allCatNames[i]
+              if (!name) return c
+              return !hiddenCategories.has(name) ? HIDDEN_COLOR : c
+            })
+          }
         } else {
           categoryColName = categoryColumn
           colors = CATEGORY_COLORS.slice(0, Math.max(colInfo.nUnique, 10))
@@ -291,7 +303,7 @@ export default function EmbeddingView({ brush, categoryColumn, categoryColumns,
       selection: null,
       tooltip: null,
     })
-  }, [categoryColumn, categoryColumns])
+  }, [categoryColumn, categoryColumns, hiddenCategories])
 
   // Handle resize
   useEffect(() => {

bionemo-recipes/interpretability/sparse_autoencoders/recipes/codonfm/codon_dashboard/src/EmbeddingView.jsx

@@ -349,6 +349,24 @@ const FeatureCard = forwardRef(function FeatureCard({ feature, isHighlighted, fo
       lines.push('')
     }
 
+    // GSEA enrichment section
+    const gseaCsvFields = [
+      { key: 'gsea_overall_best', label: 'GSEA Overall Best' },
+      { key: 'gsea_GO_Biological_Process', label: 'GSEA GO Biological Process' },
+      { key: 'gsea_GO_Molecular_Function', label: 'GSEA GO Molecular Function' },
+      { key: 'gsea_GO_Cellular_Component', label: 'GSEA GO Cellular Component' },
+      { key: 'gsea_InterPro_Domains', label: 'GSEA InterPro Domains' },
+      { key: 'gsea_GO_Slim', label: 'GSEA GO Slim' },
+    ]
+    const gseaLines = gseaCsvFields
+      .filter(({ key }) => feature[key] && feature[key] !== 'unlabeled')
+      .map(({ key, label }) => `${label},${feature[key]}`)
+    if (gseaLines.length > 0) {
+      lines.push('=== GSEA ENRICHMENT ===')
+      gseaLines.forEach(l => lines.push(l))
+      lines.push('')
+    }
+
     // Examples section
     if (examples && examples.length > 0) {
       lines.push('=== ACTIVATION EXAMPLES ===')
@@ -657,6 +675,26 @@ const FeatureCard = forwardRef(function FeatureCard({ feature, isHighlighted, fo
             if (ann.cpg) tags.push({ label: `CpG enriched`, color: '#fce4ec' })
             if (ann.position) tags.push({ label: `N-terminal`, color: '#e8f5e9' })
 
+            // GSEA enrichment tags
+            const gseaFields = [
+              { key: 'gsea_GO_Biological_Process', prefix: 'GO:BP', color: '#e8eaf6' },
+              { key: 'gsea_GO_Molecular_Function', prefix: 'GO:MF', color: '#ede7f6' },
+              { key: 'gsea_GO_Cellular_Component', prefix: 'GO:CC', color: '#e0f2f1' },
+              { key: 'gsea_InterPro_Domains', prefix: 'InterPro', color: '#fff8e1' },
+              { key: 'gsea_GO_Slim', prefix: 'GO Slim', color: '#f1f8e9' },
+            ]
+            for (const { key, prefix, color } of gseaFields) {
+              const val = feature[key]
+              if (val && val !== 'unlabeled' && val !== 'other') {
+                tags.push({ label: `${prefix}: ${val}`, color })
+              }
+            }
+
+            // Codon optimality metrics from annotations
+            if (ann.cai != null) tags.push({ label: `CAI: ${ann.cai.toFixed(3)}`, color: '#e0f7fa' })
+            if (ann.tai != null) tags.push({ label: `tAI: ${ann.tai.toFixed(3)}`, color: '#e0f7fa' })
+            if (ann.rscu != null) tags.push({ label: `RSCU: ${ann.rscu.toFixed(2)}`, color: '#e0f7fa' })
+
             if (tags.length === 0) return null
             return (
               <div style={{ display: 'flex', flexWrap: 'wrap', gap: '4px', marginBottom: '10px' }}>

bionemo-recipes/interpretability/sparse_autoencoders/recipes/codonfm/codon_dashboard/src/FeatureCard.jsx

@@ -369,6 +369,53 @@ export default function FeatureDetailPage({ feature, examples, vocabLogits, feat
           <VocabLogitChart logits={logits} />
         </div>
 
+        {/* Gene-Level GSEA Enrichment */}
+        {(() => {
+          const gseaFields = [
+            { key: 'gsea_GO_Biological_Process', label: 'GO Biological Process' },
+            { key: 'gsea_GO_Molecular_Function', label: 'GO Molecular Function' },
+            { key: 'gsea_GO_Cellular_Component', label: 'GO Cellular Component' },
+            { key: 'gsea_InterPro_Domains', label: 'InterPro Domains' },
+            { key: 'gsea_GO_Slim', label: 'GO Slim' },
+          ]
+          const gseaEntries = gseaFields
+            .map(({ key, label }) => ({ label, value: feature[key] }))
+            .filter(e => e.value && e.value !== 'unlabeled' && e.value !== 'other')
+          const overallBest = feature.gsea_overall_best
+          if (gseaEntries.length === 0 && (!overallBest || overallBest === 'unlabeled')) return null
+          return (
+            <div style={styles.section}>
+              <div style={styles.sectionTitle}>Gene-Level Enrichment (GSEA)</div>
+              <div style={styles.sectionSubtitle}>
+                Genes ranked by activation strength, tested against GO and InterPro databases.
+              </div>
+              {overallBest && overallBest !== 'unlabeled' && (
+                <div style={{
+                  padding: '8px 12px', marginBottom: '8px', borderRadius: '6px',
+                  background: 'var(--bg-card-expanded)', border: '1px solid var(--accent)',
+                  fontSize: '13px', fontWeight: '600', color: 'var(--text-heading)',
+                }}>
+                  Best: {overallBest}
+                </div>
+              )}
+              <div style={{ display: 'grid', gridTemplateColumns: '1fr 1fr', gap: '6px' }}>
+                {gseaEntries.map(({ label, value }) => (
+                  <div key={label} style={{
+                    padding: '6px 10px', borderRadius: '4px',
+                    background: 'var(--bg-card)', border: '1px solid var(--border-card)',
+                    fontSize: '11px',
+                  }}>
+                    <div style={{ color: 'var(--text-muted)', fontSize: '9px', fontWeight: '600', marginBottom: '2px' }}>
+                      {label}
+                    </div>
+                    <div style={{ color: 'var(--text-primary)' }}>{value}</div>
+                  </div>
+                ))}
+              </div>
+            </div>
+          )
+        })()}
+
         {/* Codon Annotations */}
         <div style={styles.section}>
           <div style={styles.sectionTitle}>Codon-Level Annotations</div>

bionemo-recipes/interpretability/sparse_autoencoders/recipes/codonfm/codon_dashboard/src/FeatureDetailPage.jsx

@@ -133,6 +133,14 @@ def run_train(cfg: DictConfig, cache_dir: Path, output_dir: Path) -> None:  # no
         cmd.append("--normalize-input")
     if t.get("max_grad_norm"):
         cmd.extend(["--max-grad-norm", str(t.max_grad_norm)])
+    if t.get("lr_schedule", "constant") != "constant":
+        cmd.extend(["--lr-schedule", str(t.lr_schedule)])
+    if t.get("lr_min", 0.0) != 0.0:
+        cmd.extend(["--lr-min", str(t.lr_min)])
+    if t.get("lr_decay_steps"):
+        cmd.extend(["--lr-decay-steps", str(t.lr_decay_steps)])
+    if t.get("warmup_steps", 0) > 0:
+        cmd.extend(["--warmup-steps", str(t.warmup_steps)])
 
     if t.wandb_enabled:
         cmd.append("--wandb")

bionemo-recipes/interpretability/sparse_autoencoders/recipes/codonfm/run.py

@@ -42,6 +42,10 @@ train:
   wandb_enabled: false
   wandb_project: sae_codonfm_recipe
   max_grad_norm: null
+  lr_schedule: constant
+  lr_min: 0.0
+  lr_decay_steps: null
+  warmup_steps: 0
 
 # Eval
 eval: