From d36dbe0b2bedb2a40fbe9104fd156fc2636fe9bf Mon Sep 17 00:00:00 2001
From: w-ahern <bwahern@gmail.com>
Date: Fri, 12 Sep 2025 00:23:19 -0700
Subject: [PATCH] Fix pymol color palette by removing reliance on pymolrc

---
 rf_diffusion/dev/analyze.py      | 5 +++--
 rf_diffusion/dev/paper_vis.py    | 3 ++-
 rf_diffusion/dev/pymol.py        | 9 ++++++++-
 rf_diffusion/dev/show_bench.py   | 3 ++-
 rf_diffusion/dev/show_tip_row.py | 3 ++-
 5 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/rf_diffusion/dev/analyze.py b/rf_diffusion/dev/analyze.py
index 57cb495..a3996f4 100644
--- a/rf_diffusion/dev/analyze.py
+++ b/rf_diffusion/dev/analyze.py
@@ -21,6 +21,7 @@
 REPO_DIR = os.path.dirname(PKG_DIR)
 
 from rf_diffusion.dev.pymol import cmd
+from rf_diffusion.dev.pymol import mass_paper_rainbow_sel
 
 import rf_diffusion.estimate_likelihood as el
 from rf_diffusion.inference import utils
@@ -1189,7 +1190,7 @@ def show_paper_pocket(row, des=True, ligand=False):
         #cmd.align(des.name, af2_scaffold.name)
         cmd.align(f'{des.motif_sele()} and name ca', f'{native.motif_sele()} and name ca')
         cmd.show('cartoon', f'{des.name}')
-        cmd.do(f'mass_paper_rainbow_sel {des.name}')
+        mass_paper_rainbow_sel(des.name)
 
     # AF2 sidechains
     cmd.show('licorice', f'({af2.motif_sele()}) and not (name o)')
@@ -1268,7 +1269,7 @@ def show_paper_pocket_af2(row, b=None, des=True, ligand=False, traj_types=None,
         if has_motif:
             cmd.align(f'{des.motif_sele()} and name ca', f'{native.motif_sele()} and name ca')
         cmd.show('cartoon', f'{des.name}')
-        cmd.do(f'mass_paper_rainbow_sel {des.name}')
+        mass_paper_rainbow_sel(des.name)
 
     if show_af2 and has_motif:
         # AF2 sidechains
diff --git a/rf_diffusion/dev/paper_vis.py b/rf_diffusion/dev/paper_vis.py
index 53a220a..2eea4cb 100644
--- a/rf_diffusion/dev/paper_vis.py
+++ b/rf_diffusion/dev/paper_vis.py
@@ -9,6 +9,7 @@
 from rf_diffusion.dev import show_bench
 from rf_diffusion.dev import analyze
 from rf_diffusion.dev import pymol_atom_colors
+from rf_diffusion.dev.pymol import mass_paper_rainbow_sel
 
 cmd = analyze.cmd
 
@@ -80,7 +81,7 @@ def make_transparent_rainbow(
     diffused_transparency = 0.45):
     cmd.set('cartoon_transparency', diffused_transparency)
     cmd.set('stick_transparency', diffused_transparency, e['sidechains_diffused'])
-    cmd.do(f'mass_paper_rainbow {e["protein"]}')
+    mass_paper_rainbow_sel(f'{e["protein"]}')
 
 def show_backbone_cartoon(e):
     cmd.show_as('cartoon', f"({e['protein']}) and not (({e['sidechains_motif']}) or ({e['sidechains_diffused']}))")
diff --git a/rf_diffusion/dev/pymol.py b/rf_diffusion/dev/pymol.py
index 4523b26..d538f27 100644
--- a/rf_diffusion/dev/pymol.py
+++ b/rf_diffusion/dev/pymol.py
@@ -24,7 +24,7 @@ def __call__(self, *args, **kw):
             raise Exception(f"cmd.{self.name}('{','.join(all_args)})'") from e
 
 def get_cmd(pymol_url='http://localhost:9123'):
-    cmd = xmlrpclib.ServerProxy(pymol_url)
+    cmd = xmlrpclib.ServerProxy(pymol_url, allow_none=True)
     if not  ('ipd' in pymol_url or 'localhost' in pymol_url):
         make_network_cmd(cmd)
     return cmd
@@ -41,7 +41,14 @@ def init(pymol_url='http://localhost:9123', init_colors=False):
     
     if init_colors:
         set_colors()
+        cmd.set('sphere_scale', 0.3)
     
+def mass_paper_rainbow_sel(name):
+    """
+    Colors the given selection with a pastel gradient according to residue index.
+    """
+    cmd.spectrum("count", "paper_navaho paper_melon paper_pink paper_purple paper_lightblue paper_blue paper_darkblue", name, None, None, 0, 1)
+
 def set_colors():
     """
     Creates extra colors in pymol 
diff --git a/rf_diffusion/dev/show_bench.py b/rf_diffusion/dev/show_bench.py
index ab982f2..ef13c3e 100755
--- a/rf_diffusion/dev/show_bench.py
+++ b/rf_diffusion/dev/show_bench.py
@@ -17,6 +17,7 @@
 from rf_diffusion.dev import show_tip_pa
 from rf_diffusion.dev import analyze
 from rf_diffusion.parsers import parse_pdb_lines_target
+from rf_diffusion.dev.pymol import mass_paper_rainbow_sel
 
 logger = logging.getLogger(__name__)
 
@@ -458,7 +459,7 @@ def format_ppi(all_entities):
                 continue
             e.selectors['target'] = "chain B and not hetatm"
             cmd.color('paper_teal', e['target'])
-            cmd.do(f"mass_paper_rainbow_sel ({e.NOT('target')} and not hetatm)")
+            mass_paper_rainbow_sel(f"({e.NOT('target')} and not hetatm)")
 
 def pseudoatom(
         cmd,
diff --git a/rf_diffusion/dev/show_tip_row.py b/rf_diffusion/dev/show_tip_row.py
index eaa54c0..dbd5eca 100644
--- a/rf_diffusion/dev/show_tip_row.py
+++ b/rf_diffusion/dev/show_tip_row.py
@@ -5,6 +5,7 @@
 import functools
 
 from rf_diffusion.dev import analyze
+from rf_diffusion.dev.pymol import mass_paper_rainbow_sel
 cmd = analyze.cmd
 
 class PymolPalette:
@@ -176,7 +177,7 @@ def color_selectors(selectors, carbon=True, verbose=False, des_color=None, hetat
             print(f'{color} --> {sel}')
         if des_color and j==0:
             if des_color == 'rainbow':
-                cmd.do(f'mass_paper_rainbow_sel {sel}')
+                mass_paper_rainbow_sel(sel)
             else:
                 cmd.color(des_color, sel)
         else: