feat: Implement smooth custom color transitions using NGL custom schemes

Author: AmirMohammad Cheraghali

src/components/ProteinViewer.tsx

@@ -136,6 +136,7 @@ export const ProteinViewer = forwardRef<ProteinViewerRef, ProteinViewerProps>(({
     const highlightComponentRef = useRef<any>(null);
     const isMounted = useRef(true);
     const onHoverRef = useRef(onHover);
+    const viewerInstanceId = useRef(crypto.randomUUID()); // Unique ID for custom schemes
 
     // Update ref when prop changes
     useEffect(() => {
@@ -1761,6 +1762,16 @@ export const ProteinViewer = forwardRef<ProteinViewerRef, ProteinViewerProps>(({
 
             console.log("Coloring Debug:", { currentColoring, repType, hasValidCustomRules, rules: customColors });
 
+            // Capture the base coloring intent (e.g. 'chainid') to use as fallback in our custom scheme
+            const baseColoringForScheme = currentColoring;
+
+            // If we have custom rules, we want to bypass the manual 'chainid' and 'charge' blocks below
+            // and fall through to the generic 'else' block where our new Custom Scheme logic lives.
+            // We force currentColoring to a value that won't match the specific if-blocks.
+            if (hasValidCustomRules) {
+                currentColoring = 'custom_override' as any;
+            }
+
             // --- STRATEGY: MULTI-REPRESENTATION OVERLAY (RESTORED & IMPROVED) ---
             // NGL Custom Schemes proved fragile for this user.
             // We implementation "High Contrast Chain Coloring" by EXPLICITLY adding a representation for each chain.
@@ -1857,7 +1868,58 @@ export const ProteinViewer = forwardRef<ProteinViewerRef, ProteinViewerProps>(({
                     return PALETTES[p] || undefined;
                 };
 
-                // Unified cartoon with optimized parameters for arrows and helices
+                // 2. Custom Color Scheme with Smooth Blending
+                // Instead of overlays (which cause jagged geometry), we create a custom NGL scheme
+                // that delegates to the base scheme for non-overridden atoms.
+
+                let finalColoring = currentColoring;
+
+                if (hasValidCustomRules) {
+                    const schemeId = `custom_scheme_${viewerInstanceId.current}`;
+                    const atomColorMap = new Map<number, number>();
+
+                    // Pre-calculate atom colors
+                    customColors.forEach(rule => {
+                        if (rule.color && rule.target) {
+                            const colorHex = new window.NGL.Color(rule.color).getHex();
+                            try {
+                                // Efficiently populate map using NGL selection
+                                component.structure.eachAtom((atom: any) => {
+                                    atomColorMap.set(atom.index, colorHex);
+                                }, new window.NGL.Selection(rule.target));
+                            } catch (e) {
+                                console.warn("Invalid selection for custom rule:", rule.target);
+                            }
+                        }
+                    });
+
+                    // Register (or overwrite) the custom scheme
+                    // We capture 'atomColorMap' and 'currentColoring' in the closure
+                    window.NGL.ColormakerRegistry.addScheme(function (this: any, params: any) {
+                        this.parameters = params;
+
+                        // Instantiate the base scheme (fallback)
+                        // We must handle 'standard' or other aliases if necessary, but usually NGL names match
+                        // If currentColoring is not a registered scheme, fallback to 'uniform' (white)
+                        const BaseSchemeClass = window.NGL.ColormakerRegistry.getScheme(baseColoringForScheme);
+                        this.baseScheme = BaseSchemeClass ? new BaseSchemeClass(params) : null;
+
+                        this.atomColor = function (atom: any) {
+                            const custom = atomColorMap.get(atom.index);
+                            if (custom !== undefined) return custom;
+
+                            if (this.baseScheme) {
+                                return this.baseScheme.atomColor(atom);
+                            }
+                            return 0xFFFFFF; // Fallback white
+                        };
+                    }, schemeId);
+
+                    finalColoring = schemeId;
+                    console.log(`Applied smooth custom coloring: ${schemeId} with ${atomColorMap.size} atom overrides`);
+                }
+
+                // Unified cartoon with optimized parameters
                 if (repType === 'cartoon') {
                     // Force recalculation of secondary structure
                     try {
@@ -1867,10 +1929,10 @@ export const ProteinViewer = forwardRef<ProteinViewerRef, ProteinViewerProps>(({
                     } catch (e) { }
 
                     const params: any = {
-                        color: currentColoring,
-                        aspectRatio: 5,          // Flat arrows for sheets
-                        subdiv: 12,              // Smooth curves
-                        radialSegments: 20,      // Smooth helix cylinders
+                        color: finalColoring,
+                        aspectRatio: 5,
+                        subdiv: 12,
+                        radialSegments: 20,
                     };
 
                     const scale = getColorScale(colorPalette);
@@ -1881,7 +1943,7 @@ export const ProteinViewer = forwardRef<ProteinViewerRef, ProteinViewerProps>(({
                 } else {
                     // Non-cartoon representations
                     const params: any = {
-                        color: currentColoring
+                        color: finalColoring
                     };
                     const scale = getColorScale(colorPalette);
                     if (scale && scale.length > 0) {
@@ -1891,25 +1953,6 @@ export const ProteinViewer = forwardRef<ProteinViewerRef, ProteinViewerProps>(({
                 }
             }
 
-            // 2. Add Custom Representations (Overlay)
-            if (hasValidCustomRules) {
-                console.log("Applying Custom Rules as Overlays:", customColors.length);
-                customColors.forEach((rule, idx) => {
-                    if (rule.color && rule.target) {
-                        try {
-                            // Add a separate representation for this rule
-                            component.addRepresentation(repType, {
-                                color: new NGL.Color(rule.color).getHex(),
-                                sele: rule.target,
-                                name: `custom_rule_${idx}`
-                            });
-                        } catch (e) {
-                            console.warn("Failed to apply custom rule:", rule, e);
-                        }
-                    }
-                });
-            }
-
             // --- OVERLAYS ---
             const tryApply = (r: string, c: string, sele: string, params: any = {}) => {
                 try { component.addRepresentation(r, { color: c, sele: sele, ...params }); } catch (e) { }