improve distinctColors function

- Add `sort` and `sampleCountFactor` parameters
- Fix clustering issues

CHANGELOG.md

@@ -8,6 +8,9 @@ Note that since we don't clearly distinguish between a public and private interf
 
 - Fix handling of PDB files with insertion codes (#945)
 - Fix de-/saturate of colors with no hue
+- Improve `distinctColors` function
+    - Add `sort` and `sampleCountFactor` parameters
+    - Fix clustering issues
 
 ## [v3.41.0] - 2023-10-15
 

src/extensions/cellpack/color/generate.ts

@@ -45,8 +45,14 @@ export function CellPackGenerateColorTheme(ctx: ThemeDataContext, props: PD.Valu
         const palette = getPalette(size, { palette: {
             name: 'generate',
             params: {
-                hue, chroma: [30, 80], luminance: [15, 85],
-                clusteringStepCount: 50, minSampleCount: 800, maxCount: 75
+                hue,
+                chroma: [30, 80],
+                luminance: [15, 85],
+                clusteringStepCount: 50,
+                minSampleCount: 800,
+                maxCount: 75,
+                sampleCountFactor: 5,
+                sort: 'contrast'
             }
         } }, { minLabel: 'Min', maxLabel: 'Max' });
         legend = palette.legend;

src/mol-util/color/distinct.ts

@@ -21,21 +21,15 @@ export const DistinctColorsParams = {
     hue: PD.Interval([1, 360], { min: 0, max: 360, step: 1 }),
     chroma: PD.Interval([40, 70], { min: 0, max: 100, step: 1 }),
     luminance: PD.Interval([15, 85], { min: 0, max: 100, step: 1 }),
+    sort: PD.Select('contrast', PD.arrayToOptions(['none', 'contrast'] as const), { description: 'no sorting leaves colors approximately ordered by hue' }),
 
     clusteringStepCount: PD.Numeric(50, { min: 10, max: 200, step: 1 }, { isHidden: true }),
-    minSampleCount: PD.Numeric(800, { min: 100, max: 5000, step: 100 }, { isHidden: true })
+    minSampleCount: PD.Numeric(800, { min: 100, max: 5000, step: 100 }, { isHidden: true }),
+    sampleCountFactor: PD.Numeric(5, { min: 1, max: 100, step: 1 }, { isHidden: true }),
 };
 export type DistinctColorsParams = typeof DistinctColorsParams
 export type DistinctColorsProps = PD.Values<typeof DistinctColorsParams>
 
-function distance(colorA: Lab, colorB: Lab) {
-    return Math.sqrt(
-        Math.pow(Math.abs(colorA[0] - colorB[0]), 2) +
-        Math.pow(Math.abs(colorA[1] - colorB[1]), 2) +
-        Math.pow(Math.abs(colorA[2] - colorB[2]), 2)
-    );
-}
-
 const LabTolerance = 2;
 const tmpCheckColorHcl = [0, 0, 0] as unknown as Hcl;
 const tmpCheckColorLab = [0, 0, 0] as unknown as Lab;
@@ -66,9 +60,9 @@ function sortByContrast(colors: Lab[]) {
     while (unsortedColors.length > 0) {
         const lastColor = sortedColors[sortedColors.length - 1];
         let nearest = 0;
-        let maxDist = Number.MIN_SAFE_INTEGER;
+        let maxDist = Number.NEGATIVE_INFINITY;
         for (let i = 0; i < unsortedColors.length; ++i) {
-            const dist = distance(lastColor, unsortedColors[i]);
+            const dist = Lab.distance(lastColor, unsortedColors[i]);
             if (dist > maxDist) {
                 maxDist = dist;
                 nearest = i;
@@ -79,18 +73,19 @@ function sortByContrast(colors: Lab[]) {
     return sortedColors;
 }
 
-function getSamples(count: number, p: DistinctColorsProps) {
-    const samples = new Map<string, Lab>();
-    const rangeDivider = Math.cbrt(count) * 1.001;
+function getSamples(count: number, p: DistinctColorsProps): Lab[] {
+    const samples = new Map<number, Lab>();
+    const rangeDivider = Math.ceil(Math.cbrt(count));
+    const hcl = Hcl();
 
     const hStep = Math.max((p.hue[1] - p.hue[0]) / rangeDivider, 1);
     const cStep = Math.max((p.chroma[1] - p.chroma[0]) / rangeDivider, 1);
     const lStep = Math.max((p.luminance[1] - p.luminance[0]) / rangeDivider, 1);
-    for (let h = p.hue[0]; h <= p.hue[1]; h += hStep) {
-        for (let c = p.chroma[0]; c <= p.chroma[1]; c += cStep) {
-            for (let l = p.luminance[0]; l <= p.luminance[1]; l += lStep) {
-                const lab = Lab.fromHcl(Lab(), Hcl.create(h, c, l));
-                if (checkColor(lab, p)) samples.set(lab.toString(), lab);
+    for (let h = p.hue[0] + hStep / 2; h <= p.hue[1]; h += hStep) {
+        for (let c = p.chroma[0] + cStep / 2; c <= p.chroma[1]; c += cStep) {
+            for (let l = p.luminance[0] + lStep / 2; l <= p.luminance[1]; l += lStep) {
+                const lab = Lab.fromHcl(Lab(), Hcl.set(hcl, h, c, l));
+                if (checkColor(lab, p)) samples.set(Lab.toColor(lab), lab);
             }
         }
     }
@@ -98,22 +93,36 @@ function getSamples(count: number, p: DistinctColorsProps) {
     return Array.from(samples.values());
 }
 
+function getClosestIndex(colors: Lab[], color: Lab) {
+    let minDist = Infinity;
+    let nearest = 0;
+
+    for (let j = 0; j < colors.length; j++) {
+        const dist = Lab.distance(color, colors[j]);
+        if (dist < minDist) {
+            minDist = dist;
+            nearest = j;
+        }
+    }
+
+    return nearest;
+}
+
 /**
  * Create a list of visually distinct colors
  */
 export function distinctColors(count: number, props: Partial<DistinctColorsProps> = {}): Color[] {
     const p = { ...PD.getDefaultValues(DistinctColorsParams), ...props };
-
     if (count <= 0) return [];
 
-    const samples = getSamples(Math.max(p.minSampleCount, count * 5), p);
+    const samples = getSamples(Math.max(p.minSampleCount, count * p.sampleCountFactor), p);
     if (samples.length < count) {
         console.warn('Not enough samples to generate distinct colors, increase sample count.');
         return (new Array(count)).fill(ColorNames.lightgrey);
     }
 
     const colors: Lab[] = [];
-    const zonesProto: (Lab[])[] = [];
+    const zonesProto: Lab[][] = [];
     const sliceSize = Math.floor(samples.length / count);
 
     for (let i = 0; i < samples.length; i += sliceSize) {
@@ -124,18 +133,17 @@ export function distinctColors(count: number, props: Partial<DistinctColorsProps
 
     for (let step = 1; step <= p.clusteringStepCount; ++step) {
         const zones = deepClone(zonesProto);
+        const sampleList = deepClone(samples); // Immediately add the closest sample for each color
 
         // Find closest color for each sample
-        for (let i = 0; i < samples.length; ++i) {
-            let minDist = Number.MAX_SAFE_INTEGER;
-            let nearest = 0;
-            for (let j = 0; j < colors.length; j++) {
-                const dist = distance(samples[i], colors[j]);
-                if (dist < minDist) {
-                    minDist = dist;
-                    nearest = j;
-                }
-            }
+        for (let i = 0; i < colors.length; ++i) {
+            const idx = getClosestIndex(sampleList, colors[i]);
+            zones[i].push(samples[idx]);
+            sampleList.splice(idx, 1);
+        }
+
+        for (let i = 0; i < sampleList.length; ++i) {
+            const nearest = getClosestIndex(colors, samples[i]);
             zones[nearest].push(samples[i]);
         }
 
@@ -144,6 +152,8 @@ export function distinctColors(count: number, props: Partial<DistinctColorsProps
         for (let i = 0; i < zones.length; ++i) {
             const zone = zones[i];
             const size = zone.length;
+            if (size === 0) continue;
+
             const Ls: number[] = [];
             const As: number[] = [];
             const Bs: number[] = [];
@@ -164,5 +174,6 @@ export function distinctColors(count: number, props: Partial<DistinctColorsProps
         if (deepEqual(lastColors, colors)) break;
     }
 
-    return sortByContrast(colors).map(c => Lab.toColor(c));
+    const sorted = p.sort === 'contrast' ? sortByContrast(colors) : colors;
+    return sorted.map(c => Lab.toColor(c));
 }

src/mol-util/color/spaces/hcl.ts

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2019-2023 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  *
@@ -43,6 +43,17 @@ namespace Hcl {
         return out;
     }
 
+    export function set(out: Hcl, h: number, c: number, l: number): Hcl {
+        out[0] = h;
+        out[1] = c;
+        out[2] = l;
+        return out;
+    }
+
+    export function hasHue(a: Hcl) {
+        return !isNaN(a[0]);
+    }
+
     const tmpFromColorLab = [0, 0, 0] as unknown as Lab;
     export function fromColor(out: Hcl, color: Color): Hcl {
         return Lab.toHcl(out, Lab.fromColor(tmpFromColorLab, color));

src/mol-util/color/spaces/lab.ts

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2019-2023 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  *
@@ -44,6 +44,21 @@ namespace Lab {
         return out;
     }
 
+    export function set(out: Lab, l: number, a: number, b: number): Lab {
+        out[0] = l;
+        out[1] = a;
+        out[2] = b;
+        return out;
+    }
+
+    /** simple eucledian distance, not perceptually uniform */
+    export function distance(a: Lab, b: Lab) {
+        const x = b[0] - a[0],
+            y = b[1] - a[1],
+            z = b[2] - a[2];
+        return Math.sqrt(x * x + y * y + z * z);
+    }
+
     export function fromColor(out: Lab, color: Color): Lab {
         const [r, g, b] = Color.toRgb(color);
         const [x, y, z] = rgbToXyz(r, g, b);