Import better m21p Chord root() functions

src/chord.ts

@@ -444,64 +444,132 @@ export class Chord extends note.NotRest {
     }
 
     /**
-     * Finds the Root of the chord, or sets it as an override.
+     * Returns the root of the chord, or sets it as an override.
+     *
+     * If `newRoot` is given, sets the root to the pitch in the chord matching it
+     * (by identity, then nameWithOctave, then name), or to `newRoot` itself
+     * if none matches, as for an implied root.
+     *
+     * `find` has three states, so it has no default:
+     * - omitted: returns the set root, else the cached one, else finds it.
+     * - true: clears any set root and finds it again.
+     * - false: returns the set root or undefined; never finds it.
+     *
+     * Throws if the chord has no pitches.
+     *
+     * @example
+     * const cSus4 = new music21.chord.Chord('C4 F4 G4');
+     * cSus4.root().nameWithOctave;  // an F9 chord in 2nd inversion
+     * // 'F4'
+     * cSus4.root('C4');
+     * cSus4.root().nameWithOctave;
+     * // 'C4'
+     * cSus4.root(undefined, { find: true }).nameWithOctave;
+     * // 'F4'
      */
-    root(newroot?: pitch.Pitch): pitch.Pitch {
-        if (newroot !== undefined) {
-            this._overrides.root = newroot;
-            this._cache.root = newroot;
+    root(
+        newRoot?: pitch.Pitch|note.Note|string,
+        { find }: { find?: boolean } = {}
+    ): pitch.Pitch|undefined {
+        if (newRoot) {
+            let newRootPitch: pitch.Pitch;
+            if (typeof newRoot === 'string') {
+                newRootPitch = new pitch.Pitch(newRoot);
+            } else if (newRoot instanceof note.Note) {
+                newRootPitch = newRoot.pitch;
+            } else {
+                newRootPitch = newRoot;
+            }
+            const pitches = this.pitches;
+            newRootPitch = pitches.find(p => p === newRootPitch)
+                ?? pitches.find(p => p.nameWithOctave === newRootPitch.nameWithOctave)
+                ?? pitches.find(p => p.name === newRootPitch.name)
+                ?? newRootPitch;
+            this._overrides.root = newRootPitch;
+            this._cache.root = newRootPitch;
             this._cache.inversion = undefined;
         }
 
+        if (find === true) {
+            this._overrides.root = undefined;
+            this._cache.inversion = undefined;
+            this._cache.root = this._findRoot();
+            return this._cache.root;
+        }
         if (this._overrides.root !== undefined) {
             return this._overrides.root;
         }
-
-        if (this._cache.root !== undefined) {
-            return this._cache.root;
+        if (find === false) {
+            return undefined;
         }
-
-        const closedChord = this.removeDuplicatePitches();
-        /* var chordBass = closedChord.bass(); */
-        const closedPitches = closedChord.pitches;
-        if (closedPitches.length === 0) {
-            throw new Music21Exception('No notes in Chord!');
-        } else if (closedPitches.length === 1) {
-            return this.pitches[0];
+        if (this._cache.root === undefined) {
+            this._cache.root = this._findRoot();
         }
-        // const indexOfPitchesWithPerfectlyStackedThirds = [];
-        const testSteps = [3, 5, 7, 2, 4, 6];
-        for (let i = 0; i < closedPitches.length; i++) {
-            const p = closedPitches[i];
-            const currentListOfThirds = [];
-            for (let tsIndex = 0; tsIndex < testSteps.length; tsIndex++) {
-                const chordStepPitch = closedChord.getChordStep(
-                    testSteps[tsIndex],
-                    p
-                );
-                if (chordStepPitch !== undefined) {
-                    // console.log(p.name + " " + testSteps[tsIndex].toString() + " " + chordStepPitch.name);
-                    currentListOfThirds.push(true);
-                } else {
-                    currentListOfThirds.push(false);
+        return this._cache.root;
+    }
+
+    /**
+     * Finds the root, usually the pitch with the most thirds stacked above it.
+     * Use `.root()` instead, which caches the result.
+     *
+     * AI-assisted (Claude) port of music21p Chord._findRoot.
+     */
+    protected _findRoot(): pitch.Pitch {
+        // score for how likely a pitch is to be a root, given whether it has
+        // a 3rd, 5th, 7th, 9th, 11th, and 13th above it.
+        const rootnessFunction = (rootThirdList: boolean[]): number => {
+            let score = 0;
+            rootThirdList.forEach((val, rootIndex) => {
+                if (val) {
+                    score += 1 / (rootIndex + 6);
                 }
-            }
-            // console.log(currentListOfThirds);
-            let hasFalse = false;
-            for (let j = 0; j < closedPitches.length - 1; j++) {
-                if (currentListOfThirds[j] === false) {
-                    hasFalse = true;
+            });
+            return score;
+        };
+
+        const nonDuplicatingPitches = this.removeDuplicatePitches().pitches;
+        const lenPitches = nonDuplicatingPitches.length;
+        if (!lenPitches) {
+            throw new Music21Exception(`no pitches in chord ${this.stringInfo()}`);
+        }
+        if (lenPitches === 1) {
+            return this.pitches[0];
+        } else if (lenPitches === 7) {  // 13th chord
+            return this.bass();
+        }
+
+        // FIND ROOT FAST -- if one pitch has perfectly stacked thirds
+        // above it, like E C G (but not C E B-), return it.
+        const stepNumsToPitches = new Map<number, pitch.Pitch>();
+        for (const p of nonDuplicatingPitches) {
+            stepNumsToPitches.set(pitch.nameToSteps[p.step], p);
+        }
+        const stepNums = [...stepNumsToPitches.keys()].sort((a, b) => a - b);
+        for (let startIndex = 0; startIndex < lenPitches; startIndex++) {
+            let allAreThirds = true;
+            let lastStepNum = stepNums[startIndex];
+            for (let endIndex = startIndex + 1; endIndex < startIndex + lenPitches; endIndex++) {
+                const endStepNum = stepNums[endIndex % lenPitches];
+                if (![2, -5].includes(endStepNum - lastStepNum)) {
+                    allAreThirds = false;
+                    break;
                 }
+                lastStepNum = endStepNum;
             }
-            if (hasFalse === false) {
-                // indexOfPitchesWithPerfectlyStackedThirds.push(i);
-                return closedChord.pitches[i]; // should do more, but fine...
-                // should test rootedness function, etc. 13ths. etc.
+            if (allAreThirds) {
+                return stepNumsToPitches.get(stepNums[startIndex]);
             }
         }
-        const newRoot = closedChord.pitches[0]; // fallback, just return the bass...
-        this._cache.root = newRoot;
-        return newRoot;
+
+        // FIND ROOT SLOW -- return the highest scoring pitch
+        const orderedChordSteps = [3, 5, 7, 2, 4, 6];
+        const rootnessScores = nonDuplicatingPitches.map(p => {
+            const thisStepNum = pitch.nameToSteps[p.step];
+            return rootnessFunction(orderedChordSteps.map(
+                chordStepTest => stepNumsToPitches.has((thisStepNum + chordStepTest - 1) % 7)
+            ));
+        });
+        return nonDuplicatingPitches[rootnessScores.indexOf(Math.max(...rootnessScores))];
     }
 
     /**
@@ -538,8 +606,10 @@ export class Chord extends note.NotRest {
      * (by identity, then nameWithOctave, then name).  If no pitch matches,
      * throws unless `allowAdd` is true, in which case the pitch is added.
      *
-     * If `find` is false, returns undefined unless the bass has been set.
-     * If `find` is true, clears any set bass and finds it again.
+     * `find` has three states, so it has no default:
+     * - omitted: returns the set bass, else the cached one, else finds it.
+     * - true: clears any set bass and finds it again.
+     * - false: returns the set bass or undefined; never finds it.
      *
      * return bass pitch or undefined
      *

tests/moduleTests/chord.ts

@@ -275,4 +275,61 @@ export default function tests() {
         assert.notOk(new music21.chord.Chord('C D E F G B').isSeventh());
         assert.notOk(new music21.chord.Chord().isSeventh());
     });
+
+    test('music21.chord.Chord.root', assert => {
+        const Chord = music21.chord.Chord;
+        const rootName = (pitches: string) => new Chord(pitches).root().nameWithOctave;
+
+        assert.equal(rootName('E3 C4 G5'), 'C4');
+        assert.equal(rootName('A3 E-4 G4'), 'A3', 'A dim7 no 3rd, not E-11');
+        assert.equal(rootName('E-3 A4 G4'), 'A4');
+        assert.equal(rootName('F3 A3 C4 E-4 G-4 B4 D5'), 'F3', '13th chord root is bass');
+        assert.equal(rootName('C4 F4 G4'), 'F4', 'F9 chord in 2nd inversion');
+        assert.equal(rootName('B3 D4 F4 A-4'), 'B3');
+
+        const lotsOfNotes = new Chord('E3 C4 G4 B-4 E5 G5');
+        assert.strictEqual(lotsOfNotes.root(), lotsOfNotes.pitches[1]);
+
+        // from m21p test_chord.testConstruction
+        const names = [
+            ['C4 E-5 B-3', 'C'],
+            ['C4 E-5 G-3', 'C'],
+            ['C4 F4 A3', 'F'],
+            ['C4 F4 A3 E-5', 'F'],
+            ['C4', 'C'],
+            ['C4 E5', 'C'],
+            ['C4 E4 G4', 'C'],
+            ['C4 E4 G4 B-4', 'C'],
+            ['E3 B-3 G4 C4', 'C'],
+            ['C4 E4 G4 B-4 D5', 'C'],
+            ['C4 E4 G4 B-4 D5 F5', 'C'],
+            ['C4 E4 G4 B-4 D5 F5 A-5', 'C'],
+            ['C4 E4 G4 B4', 'C'],
+        ];
+        for (const [pitches, expected] of names) {
+            assert.equal(new Chord(pitches).root().name, expected, pitches);
+        }
+
+        assert.equal(new Chord('G4 C5').inversion(), 2);
+
+        assert.throws(() => new Chord().root(), /no pitches in chord/);
+    });
+
+    test('music21.chord.Chord.root setting', assert => {
+        const cSus4 = new music21.chord.Chord('C4 F4 G4');
+        cSus4.root('C4');
+        assert.strictEqual(cSus4.root(), cSus4.pitches[0], 'string root snaps to chord pitch');
+        assert.equal(cSus4.root(undefined, { find: true }).nameWithOctave, 'F4');
+        assert.equal(cSus4.root().nameWithOctave, 'F4', 'find: true removed the override');
+
+        // implied root not in the chord
+        const vo9 = new music21.chord.Chord('B3 D4 F4 A-4');
+        vo9.root(new music21.pitch.Pitch('G3'));
+        assert.equal(vo9.root().nameWithOctave, 'G3');
+        assert.equal(vo9.stringInfo(), 'B3 D4 F4 A-4', 'pitches untouched');
+
+        const e = new music21.chord.Chord('E3 G3 B4');
+        assert.equal(e.root(undefined, { find: false }), undefined);
+        assert.equal(e.root().nameWithOctave, 'E3');
+    });
 }