From 85a044cdec93c4eedf3f39d340de41ff90053862 Mon Sep 17 00:00:00 2001
From: Janosh Riebesell <janosh.riebesell@gmail.com>
Date: Mon, 29 Jan 2024 10:27:59 +0100
Subject: [PATCH] define type AtomsGraph and subtypes for
 (Structure|Molecule)Graph

rename symmetrize_structure() to get_pbc_image_sites()

add number_electrons+electron_label_props to BohrAtom.svelte
---
 package.json                            | 30 +++++++-------
 src/lib/BohrAtom.svelte                 | 28 ++++++++++++-
 src/lib/material/index.ts               |  1 +
 src/lib/structure/Structure.svelte      |  8 ++--
 src/lib/structure/StructureScene.svelte |  4 +-
 src/lib/structure/index.ts              | 55 +++++++++++++++++++++----
 tests/unit/structure-utils.test.ts      | 14 +++----
 7 files changed, 102 insertions(+), 38 deletions(-)

diff --git a/package.json b/package.json
index 7c1ce5c..40ad6b9 100644
--- a/package.json
+++ b/package.json
@@ -24,9 +24,9 @@
   },
   "dependencies": {
     "@iconify/svelte": "^3.1.6",
-    "@sveltejs/kit": "^2.3.2",
+    "@sveltejs/kit": "^2.5.0",
     "@threlte/core": "7.0.12",
-    "@threlte/extras": "8.1.2",
+    "@threlte/extras": "8.3.0",
     "d3": "^7.8.5",
     "d3-array": "^3.2.4",
     "d3-color": "^3.1.0",
@@ -36,15 +36,15 @@
     "d3-scale-chromatic": "^3.0.0",
     "d3-shape": "^3.2.0",
     "highlight.js": "^11.9.0",
-    "svelte": "4.2.8",
+    "svelte": "4.2.9",
     "svelte-multiselect": "^10.2.0",
     "svelte-zoo": "^0.4.9",
-    "three": "^0.160.0"
+    "three": "^0.160.1"
   },
   "devDependencies": {
-    "@playwright/test": "^1.40.1",
+    "@playwright/test": "^1.41.1",
     "@sveltejs/adapter-static": "3.0.1",
-    "@sveltejs/package": "^2.2.5",
+    "@sveltejs/package": "^2.2.6",
     "@sveltejs/vite-plugin-svelte": "^3.0.1",
     "@types/d3-array": "^3.2.1",
     "@types/d3-color": "^3.1.3",
@@ -53,16 +53,16 @@
     "@types/d3-scale-chromatic": "^3.0.3",
     "@types/d3-shape": "^3.1.6",
     "@types/three": "^0.160.0",
-    "@typescript-eslint/eslint-plugin": "^6.18.1",
-    "@typescript-eslint/parser": "^6.18.1",
-    "@vitest/coverage-v8": "^1.2.0",
+    "@typescript-eslint/eslint-plugin": "^6.19.1",
+    "@typescript-eslint/parser": "^6.19.1",
+    "@vitest/coverage-v8": "^1.2.1",
     "eslint": "^8.56.0",
     "eslint-plugin-svelte": "^2.35.1",
-    "hastscript": "^8.0.0",
-    "jsdom": "^23.2.0",
+    "hastscript": "^9.0.0",
+    "jsdom": "^24.0.0",
     "mdsvex": "^0.11.0",
     "mdsvexamples": "^0.4.1",
-    "prettier": "^3.2.2",
+    "prettier": "^3.2.4",
     "prettier-plugin-svelte": "^3.1.2",
     "rehype-autolink-headings": "^7.1.0",
     "rehype-katex-svelte": "^1.2.0",
@@ -71,11 +71,11 @@
     "sharp": "^0.33.2",
     "svelte-check": "^3.6.3",
     "svelte-preprocess": "^5.1.3",
-    "svelte-toc": "^0.5.6",
+    "svelte-toc": "^0.5.7",
     "svelte2tsx": "^0.7.0",
     "typescript": "5.3.3",
-    "vite": "^5.0.11",
-    "vitest": "^1.2.0"
+    "vite": "^5.0.12",
+    "vitest": "^1.2.1"
   },
   "keywords": [
     "svelte",
diff --git a/src/lib/BohrAtom.svelte b/src/lib/BohrAtom.svelte
index 7f30582..638119e 100644
--- a/src/lib/BohrAtom.svelte
+++ b/src/lib/BohrAtom.svelte
@@ -4,7 +4,7 @@
   export let name: string = `` // usually Hydrogen, Helium, etc. but can be anything
   export let shells: number[] // e.g. [2, 8, 6] for sulfur
   export let adapt_size = false
-  export let shell_width = 15 // TODO SVG is fixed so increasing this will make large atoms overflow
+  export let shell_width = 20 // TODO SVG is fixed so increasing this will make large atoms overflow
   export let size = adapt_size ? (shells.length + 1) * 2 * shell_width + 50 : 270
   export let base_fill = `white`
   export let orbital_period = 3 // time for inner-most electron orbit in seconds, 0 for no motion
@@ -14,6 +14,13 @@
   export let electron_props: Record<string, string | number> = {}
   export let highlight_shell: number | null = null
   export let style = ``
+  // if function, it'll be called with electron index and should return a string
+  export let number_electrons:
+    | boolean
+    | 'hierarchical'
+    | 'sequential'
+    | ((idx: number) => string) = false
+  export let electron_label_props: Record<string, string | number> = {}
 
   // Bohr atom electron orbital period is given by
   // T = (n^3 h^3) / (4pi^2 m K e^4 Z^2) = 1.52 * 10^-16 * n^3 / Z^2 s
@@ -77,6 +84,25 @@
         <circle class="electron" cx={elec_x} cy={elec_y} {..._electron_props}>
           <title>Electron {elec_idx + 1}</title>
         </circle>
+        {#if number_electrons}
+          <text
+            x={elec_x}
+            y={elec_y}
+            {...electron_label_props}
+            transform="rotate({(elec_idx * 360) / electrons} {elec_x} {elec_y})"
+          >
+            {#if typeof number_electrons === `function`}
+              {number_electrons(elec_idx)}
+            {:else if number_electrons === `hierarchical`}
+              {shell_idx + 1}.{elec_idx + 1}
+              <!-- {:else if [`sequential`, true].includes(number_electrons)} -->
+            {:else}
+              {@const nth_electron =
+                shells.slice(0, shell_idx).reduce((a, b) => a + b, 0) + elec_idx + 1}
+              {nth_electron}
+            {/if}
+          </text>
+        {/if}
       {/each}
     </g>
   {/each}
diff --git a/src/lib/material/index.ts b/src/lib/material/index.ts
index da428c2..696fe69 100644
--- a/src/lib/material/index.ts
+++ b/src/lib/material/index.ts
@@ -1 +1,2 @@
 export { default as MaterialCard } from './MaterialCard.svelte'
+export { default as SymmetryCard } from './SymmetryCard.svelte'
diff --git a/src/lib/structure/Structure.svelte b/src/lib/structure/Structure.svelte
index c9c32e7..e45f808 100644
--- a/src/lib/structure/Structure.svelte
+++ b/src/lib/structure/Structure.svelte
@@ -1,7 +1,7 @@
 <script lang="ts">
   import { browser } from '$app/environment'
-  import type { ElementSymbol, StructureOrMolecule, Vector } from '$lib'
-  import { alphabetical_formula, get_elem_amounts, symmetrize_structure } from '$lib'
+  import type { Atoms, ElementSymbol, Vector } from '$lib'
+  import { alphabetical_formula, get_elem_amounts, get_pbc_image_sites } from '$lib'
   import { download } from '$lib/api'
   import { element_color_schemes } from '$lib/colors'
   import { element_colors } from '$lib/stores'
@@ -11,7 +11,7 @@
   import StructureScene from './StructureScene.svelte'
 
   // output of pymatgen.core.Structure.as_dict()
-  export let structure: StructureOrMolecule | undefined = undefined
+  export let structure: Atoms | undefined = undefined
   // scale factor for atomic radii
   export let atom_radius: number = 0.5
   // whether to use the same radius for all atoms. if not, the radius will be
@@ -403,7 +403,7 @@
 
     <Canvas rendererParameters={{ preserveDrawingBuffer: true }}>
       <StructureScene
-        structure={show_image_atoms ? symmetrize_structure(structure) : structure}
+        structure={show_image_atoms ? get_pbc_image_sites(structure) : structure}
         {show_atoms}
         {show_bonds}
         {show_cell}
diff --git a/src/lib/structure/StructureScene.svelte b/src/lib/structure/StructureScene.svelte
index cc7cf8b..2b45f16 100644
--- a/src/lib/structure/StructureScene.svelte
+++ b/src/lib/structure/StructureScene.svelte
@@ -1,5 +1,5 @@
 <script lang="ts">
-  import type { BondPair, Site, StructureOrMolecule, Vector } from '$lib'
+  import type { Atoms, BondPair, Site, Vector } from '$lib'
   import {
     Bond,
     Lattice,
@@ -21,7 +21,7 @@
   import * as bonding_strategies from './bonding'
 
   // output of pymatgen.core.Structure.as_dict()
-  export let structure: StructureOrMolecule | undefined = undefined
+  export let structure: Atoms | undefined = undefined
   // scale factor for atomic radii
   export let atom_radius: number = 0.5
   // whether to use the same radius for all atoms. if not, the radius will be
diff --git a/src/lib/structure/index.ts b/src/lib/structure/index.ts
index 15671ec..b486036 100644
--- a/src/lib/structure/index.ts
+++ b/src/lib/structure/index.ts
@@ -44,15 +44,44 @@ export type PymatgenStructure = {
   id?: string
 }
 
+export type Edge = {
+  to_jimage: [number, number, number]
+  id: number
+  key: number
+}
+
+export type Node = {
+  id: number
+}
+
+export type Graph = {
+  directed: boolean
+  multigraph: boolean
+  graph: [
+    ['edge_weight_name', null] | ['edge_weight_units', null] | ['name', string],
+  ]
+  nodes: Node[]
+  adjacency: Edge[][]
+}
+
+export type StructureGraph = {
+  '@module': string
+  '@class': string
+  structure: PymatgenStructure
+  graphs: Graph[]
+}
+
 // [atom_pos_1, atom_pos_2, atom_idx_1, atom_idx_2, bond_length]
 export type BondPair = [Vector, Vector, number, number, number]
 
 export type IdStructure = PymatgenStructure & { id: string }
+export type StructureWithGraph = IdStructure & { graph: Graph }
 
 // remove lattice from pymatgen Structure
 export type PymatgenMolecule = Omit<PymatgenStructure, 'lattice'>
 
-export type StructureOrMolecule = PymatgenStructure | PymatgenMolecule
+export type Atoms = PymatgenStructure | PymatgenMolecule
+export type AtomsGraph = Atoms & { graph: Graph }
 
 export function get_elem_amounts(structure: PymatgenStructure) {
   const elements: Partial<Record<ElementSymbol, number>> = {}
@@ -111,23 +140,29 @@ export function density(structure: PymatgenStructure, prec = `.2f`) {
 }
 
 function generate_permutations(length: number): number[][] {
+  // generate all permutations of 0s and 1s of length `length`
   const result: number[][] = []
-  for (let i = 0; i < Math.pow(2, length); i++) {
-    const binaryString = i.toString(2).padStart(length, `0`)
+  for (let idx = 0; idx < Math.pow(2, length); idx++) {
+    const binaryString = idx.toString(2).padStart(length, `0`)
     result.push(Array.from(binaryString).map(Number))
   }
   return result
 }
 
-// this function finds all atoms needed to make the unit cell symmetrically occupied
 export function find_image_atoms(
   structure: PymatgenStructure,
   { tolerance = 0.05 }: { tolerance?: number } = {},
   // fractional tolerance for determining if a site is at the edge of the unit cell
 ): [number, Vector][] {
+  /*
+    This function finds all atoms on corners and faces of the unit cell needed to make the cell symmetrically occupied.
+    It returns an array of [atom_idx, image_xyz] pairs where atom_idx is the index of
+    the original atom and image_xyz is the position of one of its images.
+  */
   if (!structure.lattice) return []
+
   const edge_sites: Array<[number, Vector]> = []
-  const permutations = generate_permutations(3)
+  const permutations = generate_permutations(3) //  [1, 0, 0], [0, 1, 0], etc.
   const lattice_vecs = structure.lattice?.matrix
 
   for (const [idx, site] of structure.sites.entries()) {
@@ -135,6 +170,7 @@ export function find_image_atoms(
     edge_sites.push([idx, site.xyz])
 
     // Check if the site is at the edge and determine its image
+    // based on whether fractional coordinates are close to 0 or 1
     const edges: number[] = [0, 1, 2].filter(
       (idx) =>
         Math.abs(abc[idx]) < tolerance || Math.abs(abc[idx] - 1) < tolerance,
@@ -142,13 +178,14 @@ export function find_image_atoms(
 
     if (edges.length > 0) {
       for (const perm of permutations) {
-        let img_xyz: Vector = [...site.xyz] // Make a copy of the array
+        let img_xyz: Vector = [...site.xyz] // copy site.xyz
         for (const edge of edges) {
           if (perm[edge] === 1) {
-            // Image atom at the opposite edge
             if (Math.abs(abc[edge]) < tolerance) {
+              // if fractional coordinate is close to 0, add lattice vector to get image location
               img_xyz = add(img_xyz, lattice_vecs[edge])
             } else {
+              // if fractional coordinate is close to 1, subtract lattice vector to get image location
               img_xyz = add(img_xyz, scale(lattice_vecs[edge], -1))
             }
           }
@@ -162,7 +199,7 @@ export function find_image_atoms(
 }
 
 // this function takes a pymatgen Structure and returns a new one with all the image atoms added
-export function symmetrize_structure(
+export function get_pbc_image_sites(
   ...args: Parameters<typeof find_image_atoms>
 ): PymatgenStructure {
   const edge_sites = find_image_atoms(...args)
@@ -181,7 +218,7 @@ export function symmetrize_structure(
   return symmetrized_structure
 }
 
-export function get_center_of_mass(struct_or_mol: StructureOrMolecule): Vector {
+export function get_center_of_mass(struct_or_mol: Atoms): Vector {
   let center: Vector = [0, 0, 0]
   let total_weight = 0
 
diff --git a/tests/unit/structure-utils.test.ts b/tests/unit/structure-utils.test.ts
index ef942c9..710bba6 100644
--- a/tests/unit/structure-utils.test.ts
+++ b/tests/unit/structure-utils.test.ts
@@ -1,4 +1,4 @@
-import * as module from '$lib/structure'
+import * as struct_utils from '$lib/structure'
 import { structures } from '$site'
 import fs from 'fs'
 import { describe, expect, test } from 'vitest'
@@ -52,7 +52,7 @@ describe.each(structures)(`structure-utils`, (structure) => {
   test.runIf(id in ref_data)(
     `get_elem_amount should return the correct element amounts for a given structure`,
     () => {
-      const result = module.get_elem_amounts(structure)
+      const result = struct_utils.get_elem_amounts(structure)
       expect(JSON.stringify(result), id).toBe(JSON.stringify(expected.amounts))
     },
   )
@@ -60,7 +60,7 @@ describe.each(structures)(`structure-utils`, (structure) => {
   test.runIf(id in ref_data)(
     `get_elements should return the unique elements in a given structure`,
     () => {
-      const result = module.get_elements(structure)
+      const result = struct_utils.get_elements(structure)
       expect(JSON.stringify(result), id).toBe(
         JSON.stringify(Object.keys(expected.amounts).sort()),
       )
@@ -70,14 +70,14 @@ describe.each(structures)(`structure-utils`, (structure) => {
   test.runIf(id in ref_data)(
     `density should return the correct density for a given structure`,
     () => {
-      const result = module.density(structure)
+      const result = struct_utils.density(structure)
       expect(Number(result), id).toBe(expected.density)
     },
   )
 })
 
 test.each(structures)(`find_image_atoms`, async (structure) => {
-  const image_atoms = module.find_image_atoms(structure)
+  const image_atoms = struct_utils.find_image_atoms(structure)
   // write reference data
   // fs.writeFileSync(
   //   `${__dirname}/fixtures/find_image_atoms/${structure.id}.json`,
@@ -91,7 +91,7 @@ test.each(structures)(`find_image_atoms`, async (structure) => {
 
 test.each(structures)(`symmetrize_structure`, async (structure) => {
   const orig_len = structure.sites.length
-  const symmetrized = module.symmetrize_structure(structure)
+  const symmetrized = struct_utils.get_pbc_image_sites(structure)
   const { id } = structure
   const expected = {
     'mp-1': 12,
@@ -108,7 +108,7 @@ test.each(structures)(`symmetrize_structure`, async (structure) => {
 })
 
 test.each(structures)(`get_center_of_mass for $id`, async (struct) => {
-  const center = module.get_center_of_mass(struct)
+  const center = struct_utils.get_center_of_mass(struct)
   const expected = ref_data[struct.id]?.center_of_mass
   if (!expected) return
   expect(