matterviz 0.3.4 → 0.3.6

package/dist/src/lib/trajectory/frame-reader.ts

@@ -0,0 +1,456 @@
+// Unified frame loader for XYZ and ASE trajectories (large file indexing)
+import type { ElementSymbol } from '$lib/element/types'
+import * as math from '$lib/math'
+import type {
+  FrameIndex,
+  FrameLoader,
+  ParseProgress,
+  TrajectoryFrame,
+  TrajectoryMetadata,
+} from './index'
+import { MAX_METADATA_SIZE } from './constants'
+import {
+  coerce_element_symbol,
+  convert_atomic_numbers,
+  count_xyz_frames,
+  create_trajectory_frame,
+  read_ndarray_from_view,
+  validate_3x3_matrix,
+} from './helpers'
+import { strip_compression_extensions } from './format-detect'
+
+export class TrajFrameReader implements FrameLoader {
+  private readonly format: `xyz` | `ase`
+  private global_numbers?: number[]
+
+  constructor(filename: string) {
+    const base_filename = strip_compression_extensions(filename)
+    this.format = base_filename.endsWith(`.traj`) ? `ase` : `xyz`
+  }
+
+  async get_total_frames(data: string | ArrayBuffer): Promise<number> {
+    if (this.format === `xyz`) {
+      if (data instanceof ArrayBuffer) throw new Error(`XYZ loader requires text data`)
+      return count_xyz_frames(data)
+    }
+    if (!(data instanceof ArrayBuffer)) throw new Error(`ASE loader requires binary data`)
+    const view = new DataView(data)
+    return Number(view.getBigInt64(32, true))
+  }
+
+  async build_frame_index(
+    data: string | ArrayBuffer,
+    sample_rate: number,
+    on_progress?: (progress: ParseProgress) => void,
+  ): Promise<FrameIndex[]> {
+    const total_frames = await this.get_total_frames(data)
+    const frame_index: FrameIndex[] = []
+
+    if (this.format === `xyz`) {
+      const data_str = data as string
+      const lines = data_str.trim().split(/\r?\n/)
+      const encoder = new TextEncoder()
+      const newline_sequence = data_str.includes(`\r\n`) ? `\r\n` : `\n`
+      const newline_byte_len = encoder.encode(newline_sequence).length
+
+      let [current_frame, line_idx, byte_offset] = [0, 0, 0]
+
+      while (line_idx < lines.length && current_frame < total_frames) {
+        if (!lines[line_idx]?.trim()) {
+          byte_offset += encoder.encode(lines[line_idx]).length + newline_byte_len
+          line_idx++
+          continue
+        }
+
+        const num_atoms = parseInt(lines[line_idx].trim(), 10)
+        if (isNaN(num_atoms) || num_atoms <= 0 || line_idx + num_atoms + 1 >= lines.length) {
+          byte_offset += encoder.encode(lines[line_idx]).length + newline_byte_len
+          line_idx++
+          continue
+        }
+
+        if (current_frame % sample_rate === 0) {
+          frame_index.push({
+            frame_number: current_frame,
+            byte_offset,
+            estimated_size: 0,
+          })
+        }
+
+        const frame_start = line_idx
+        line_idx += 2 + num_atoms
+        let frame_size = 0
+        for (let idx = frame_start; idx < line_idx; idx++) {
+          frame_size += encoder.encode(lines[idx]).length + newline_byte_len
+        }
+
+        if (current_frame % sample_rate === 0) {
+          frame_index[frame_index.length - 1].estimated_size = frame_size
+        }
+
+        byte_offset += frame_size
+        current_frame++
+
+        if (on_progress && current_frame % 1000 === 0) {
+          on_progress({
+            current: (current_frame / total_frames) * 100,
+            total: 100,
+            stage: `Indexing: ${current_frame}`,
+          })
+        }
+      }
+    } else {
+      const view = new DataView(data as ArrayBuffer)
+      const offsets_pos = Number(view.getBigInt64(40, true))
+
+      for (let idx = 0; idx < total_frames; idx += sample_rate) {
+        const frame_offset = Number(view.getBigInt64(offsets_pos + idx * 8, true))
+        frame_index.push({
+          frame_number: idx,
+          byte_offset: frame_offset,
+          estimated_size: 0,
+        })
+
+        if (on_progress && idx % 10000 === 0) {
+          on_progress({
+            current: (idx / total_frames) * 100,
+            total: 100,
+            stage: `Indexing ASE: ${idx}`,
+          })
+        }
+      }
+    }
+
+    return frame_index
+  }
+
+  async load_frame(
+    data: string | ArrayBuffer,
+    frame_number: number,
+  ): Promise<TrajectoryFrame | null> {
+    const actual_data_type = data instanceof ArrayBuffer ? `ArrayBuffer` : typeof data
+
+    if (this.format === `xyz`) {
+      if (typeof data !== `string`) {
+        throw new TypeError(
+          `load_frame expected string data for xyz format, received ${actual_data_type}`,
+        )
+      }
+      return this.load_xyz_frame(data, frame_number)
+    }
+    if (!(data instanceof ArrayBuffer)) {
+      throw new TypeError(
+        `load_frame expected ArrayBuffer data for ase format, received ${actual_data_type}`,
+      )
+    }
+    return this.load_ase_frame(data, frame_number)
+  }
+
+  async extract_plot_metadata(
+    data: string | ArrayBuffer,
+    options?: { sample_rate?: number; properties?: string[] },
+    on_progress?: (progress: ParseProgress) => void,
+  ): Promise<TrajectoryMetadata[]> {
+    const { sample_rate = 1, properties } = options ?? {}
+    const metadata_list: TrajectoryMetadata[] = []
+    const total_frames = await this.get_total_frames(data)
+
+    if (this.format === `xyz`) {
+      const lines = (data as string).trim().split(/\r?\n/)
+      let [current_frame, line_idx] = [0, 0]
+
+      while (line_idx < lines.length && current_frame < total_frames) {
+        if (!lines[line_idx]?.trim()) {
+          line_idx++
+          continue
+        }
+
+        const num_atoms = parseInt(lines[line_idx].trim(), 10)
+        if (isNaN(num_atoms) || num_atoms <= 0 || line_idx + num_atoms + 1 >= lines.length) {
+          line_idx++
+          continue
+        }
+
+        if (current_frame % sample_rate === 0) {
+          const comment = lines[line_idx + 1] || ``
+          let frame_metadata: TrajectoryMetadata | null = null
+          try {
+            frame_metadata = this.parse_xyz_metadata(comment, current_frame)
+          } catch (error) {
+            console.warn(
+              `Failed to parse XYZ metadata for frame ${current_frame} at line ${line_idx + 1}:`,
+              error,
+            )
+          }
+
+          if (frame_metadata && properties) {
+            const filtered = Object.fromEntries(
+              Object.entries(frame_metadata.properties).filter(([key]) =>
+                properties.includes(key),
+              ),
+            )
+            frame_metadata.properties = filtered
+          }
+
+          if (frame_metadata) metadata_list.push(frame_metadata)
+        }
+
+        line_idx += 2 + num_atoms
+        current_frame++
+
+        if (on_progress && current_frame % 5000 === 0) {
+          on_progress({
+            current: (current_frame / total_frames) * 100,
+            total: 100,
+            stage: `Extracting: ${current_frame}`,
+          })
+        }
+      }
+    } else if (this.format === `ase`) {
+      const view = new DataView(data as ArrayBuffer)
+      const n_items = Number(view.getBigInt64(32, true))
+      const offsets_pos = Number(view.getBigInt64(40, true))
+
+      for (let idx = 0; idx < n_items; idx += sample_rate) {
+        try {
+          const frame_offset = Number(view.getBigInt64(offsets_pos + idx * 8, true))
+          const json_length = Number(view.getBigInt64(frame_offset, true))
+
+          if (json_length > MAX_METADATA_SIZE) {
+            console.warn(
+              `Skipping large frame ${idx}: ${Math.round(json_length / 1024 / 1024)}MB`,
+            )
+            continue
+          }
+
+          const frame_data = JSON.parse(
+            new TextDecoder().decode(
+              new Uint8Array(data as ArrayBuffer, frame_offset + 8, json_length),
+            ),
+          )
+
+          const frame_metadata = this.parse_ase_metadata(frame_data, idx)
+
+          if (properties) {
+            const filtered = Object.fromEntries(
+              Object.entries(frame_metadata.properties).filter(([key]) =>
+                properties.includes(key),
+              ),
+            )
+            frame_metadata.properties = filtered
+          }
+
+          metadata_list.push(frame_metadata)
+
+          if (on_progress && idx % 5000 === 0) {
+            on_progress({
+              current: (idx / n_items) * 100,
+              total: 100,
+              stage: `Extracting ASE: ${idx}/${n_items}`,
+            })
+          }
+        } catch (error) {
+          console.warn(`Failed to extract metadata from ASE frame ${idx}:`, error)
+        }
+      }
+    }
+
+    return metadata_list
+  }
+
+  private load_xyz_frame(data: string, frame_number: number): TrajectoryFrame | null {
+    const lines = data.trim().split(/\r?\n/)
+    let [current_frame, line_idx] = [0, 0]
+
+    while (line_idx < lines.length && current_frame < frame_number) {
+      if (!lines[line_idx]?.trim()) {
+        line_idx++
+        continue
+      }
+      const num_atoms = parseInt(lines[line_idx].trim(), 10)
+      if (isNaN(num_atoms) || num_atoms <= 0) {
+        line_idx++
+        continue
+      }
+      line_idx += 2 + num_atoms
+      current_frame++
+    }
+
+    if (line_idx >= lines.length) return null
+    const num_atoms = parseInt(lines[line_idx].trim(), 10)
+    if (isNaN(num_atoms) || line_idx + num_atoms + 1 >= lines.length) return null
+
+    const comment = lines[line_idx + 1] || ``
+    const positions: number[][] = []
+    const elements: ElementSymbol[] = []
+
+    for (let idx = 0; idx < num_atoms; idx++) {
+      const parts = lines[line_idx + 2 + idx]?.trim().split(/\s+/)
+      if (parts?.length >= 4) {
+        const x_coord = parseFloat(parts[1])
+        const y_coord = parseFloat(parts[2])
+        const z_coord = parseFloat(parts[3])
+        if (
+          !Number.isFinite(x_coord) ||
+          !Number.isFinite(y_coord) ||
+          !Number.isFinite(z_coord)
+        ) {
+          console.warn(
+            `Skipping XYZ atom with invalid coordinates in indexed frame ${frame_number} at line ${
+              line_idx + 2 + idx
+            }`,
+          )
+          continue
+        }
+
+        const raw_symbol = parts[0]
+        const element_symbol = coerce_element_symbol(raw_symbol)
+        if (!element_symbol) {
+          console.warn(
+            `Skipping XYZ atom with unknown element symbol "${raw_symbol}" in indexed frame ${frame_number}`,
+          )
+          continue
+        }
+
+        elements.push(element_symbol)
+        positions.push([x_coord, y_coord, z_coord])
+      }
+    }
+
+    const metadata = this.parse_xyz_metadata(comment, frame_number)
+    return create_trajectory_frame(
+      positions,
+      elements,
+      undefined,
+      undefined,
+      metadata.step,
+      metadata.properties,
+    )
+  }
+
+  private load_ase_frame(data: ArrayBuffer, frame_number: number): TrajectoryFrame | null {
+    try {
+      const view = new DataView(data)
+      const n_items = Number(view.getBigInt64(32, true))
+      const offsets_pos = Number(view.getBigInt64(40, true))
+
+      if (frame_number >= n_items) return null
+
+      const frame_offset = Number(view.getBigInt64(offsets_pos + frame_number * 8, true))
+      const json_length = Number(view.getBigInt64(frame_offset, true))
+
+      const frame_data = JSON.parse(
+        new TextDecoder().decode(new Uint8Array(data, frame_offset + 8, json_length)),
+      )
+
+      const positions_ref = frame_data[`positions.`] ?? frame_data.positions
+      const positions = positions_ref?.ndarray
+        ? read_ndarray_from_view(view, positions_ref)
+        : (positions_ref as number[][])
+
+      const numbers_ref = frame_data[`numbers.`] ?? frame_data.numbers ?? this.global_numbers
+      const numbers: number[] = numbers_ref?.ndarray
+        ? read_ndarray_from_view(view, numbers_ref).flat()
+        : (numbers_ref as number[])
+
+      if (numbers) this.global_numbers = numbers
+      if (!numbers || !positions) throw new Error(`Missing atomic numbers or positions`)
+
+      const cell = frame_data.cell ? validate_3x3_matrix(frame_data.cell) : undefined
+      const metadata: Record<string, unknown> = {
+        step: frame_number,
+        ...frame_data.calculator,
+        ...frame_data.info,
+      }
+
+      if (cell) {
+        try {
+          metadata.volume = Math.abs(math.det_3x3(cell))
+        } catch (error) {
+          console.warn(`Failed to calculate volume for frame ${frame_number}:`, error)
+        }
+      }
+
+      return create_trajectory_frame(
+        positions,
+        convert_atomic_numbers(numbers),
+        cell,
+        frame_data.pbc ?? [true, true, true],
+        frame_number,
+        metadata,
+      )
+    } catch (error) {
+      console.warn(`Failed to load ASE frame ${frame_number}:`, error)
+      return null
+    }
+  }
+
+  private parse_xyz_metadata(comment: string, frame_number: number): TrajectoryMetadata {
+    const properties: Record<string, number> = {}
+
+    const patterns = {
+      energy: /(?:energy|E|etot)\s*[=:]?\s*([-+]?\d*\.?\d+(?:[eE][-+]?\d+)?)/i,
+      volume: /(?:volume|vol|V)\s*[=:]?\s*([-+]?\d*\.?\d+(?:[eE][-+]?\d+)?)/i,
+      pressure: /(?:pressure|press|P)\s*[=:]?\s*([-+]?\d*\.?\d+(?:[eE][-+]?\d+)?)/i,
+      force_max: /(?:max_force|fmax)\s*[=:]?\s*([-+]?\d*\.?\d+(?:[eE][-+]?\d+)?)/i,
+    }
+
+    Object.entries(patterns).forEach(([key, pattern]) => {
+      const match = pattern.exec(comment)
+      if (match) properties[key] = parseFloat(match[1])
+    })
+
+    const step_match = /(?:step|frame)\s*[=:]?\s*(\d+)/i.exec(comment)
+    const step = step_match ? parseInt(step_match[1]) : frame_number
+
+    return { frame_number, step, properties }
+  }
+
+  private parse_ase_metadata(
+    frame_data: Record<string, unknown>,
+    frame_number: number,
+  ): TrajectoryMetadata {
+    const properties: Record<string, number> = {}
+    const step = frame_number
+
+    if (frame_data.calculator && typeof frame_data.calculator === `object`) {
+      const calculator = frame_data.calculator as Record<string, unknown>
+      const calc_properties = [`energy`, `potential_energy`, `kinetic_energy`, `total_energy`]
+
+      for (const prop of calc_properties) {
+        if (prop in calculator && typeof calculator[prop] === `number`) {
+          properties[prop] = calculator[prop]
+        }
+      }
+    }
+
+    if (frame_data.info && typeof frame_data.info === `object`) {
+      const info = frame_data.info as Record<string, unknown>
+      const info_properties = [
+        `force_max`,
+        `force_norm`,
+        `stress_max`,
+        `stress_frobenius`,
+        `pressure`,
+        `temperature`,
+      ]
+
+      for (const prop of info_properties) {
+        if (prop in info && typeof info[prop] === `number`) {
+          properties[prop] = info[prop]
+        }
+      }
+    }
+
+    if (frame_data.cell && Array.isArray(frame_data.cell)) {
+      try {
+        const validated_cell = validate_3x3_matrix(frame_data.cell)
+        properties.volume = Math.abs(math.det_3x3(validated_cell))
+      } catch (error) {
+        console.warn(`Failed to calculate volume for ASE frame ${frame_number}:`, error)
+      }
+    }
+
+    return { frame_number, step, properties }
+  }
+}

package/dist/src/lib/trajectory/helpers.ts

@@ -0,0 +1,217 @@
+// Shared utilities for trajectory parsing
+import { ATOMIC_NUMBER_TO_SYMBOL } from '$lib/composition/parse'
+import type { ElementSymbol } from '$lib/element/types'
+import { ELEM_SYMBOLS } from '$lib/labels'
+import type { Vec3 } from '$lib/math'
+import * as math from '$lib/math'
+import type { AnyStructure } from '$lib/structure/index'
+import type { Pbc } from '$lib/structure/pbc'
+import type { TrajectoryFrame } from './index'
+
+const element_symbol_set = new Set<string>(ELEM_SYMBOLS)
+
+const is_valid_row = (row: unknown): boolean => {
+  if (!(Array.isArray(row) || (ArrayBuffer.isView(row) && `length` in row))) return false
+  return math.is_finite_vec3_like(row as ArrayLike<unknown>)
+}
+
+const is_valid_vec3 = (coords: unknown): coords is Vec3 =>
+  Array.isArray(coords) && math.is_finite_vec3_like(coords)
+
+export const is_valid_element_symbol = (symbol: string): symbol is ElementSymbol =>
+  element_symbol_set.has(symbol)
+
+export const coerce_element_symbol = (symbol: string): ElementSymbol | undefined =>
+  is_valid_element_symbol(symbol) ? symbol : undefined
+
+// Validate that data is a proper 3x3 matrix
+// Accepts both regular arrays and typed arrays (Float32Array, Float64Array, etc.)
+export function validate_3x3_matrix(data: unknown): math.Matrix3x3 {
+  if (!Array.isArray(data) || data.length !== 3) {
+    throw new Error(
+      `Expected 3x3 matrix, got array of length ${Array.isArray(data) ? data.length : `non-array`}`,
+    )
+  }
+
+  if (!data.every(is_valid_row)) {
+    throw new Error(`Invalid 3x3 matrix structure`)
+  }
+  return data as math.Matrix3x3
+}
+
+export const convert_atomic_numbers = (numbers: number[]): ElementSymbol[] =>
+  numbers.map((num) => {
+    const symbol = ATOMIC_NUMBER_TO_SYMBOL[num]
+    if (!symbol || !is_valid_element_symbol(symbol)) {
+      throw new Error(`Unknown atomic number in trajectory data: ${num}`)
+    }
+    return symbol
+  })
+
+export const create_structure = (
+  positions: number[][],
+  elements: ElementSymbol[],
+  lattice_matrix?: math.Matrix3x3,
+  pbc?: Pbc,
+  force_data?: number[][],
+): AnyStructure => {
+  if (positions.length !== elements.length) {
+    throw new Error(
+      `create_structure requires matching positions and elements lengths, got positions=${positions.length}, elements=${elements.length}`,
+    )
+  }
+  const cart_to_frac = lattice_matrix ? math.create_cart_to_frac(lattice_matrix) : null
+
+  const sites = positions.map((pos, idx) => {
+    if (!is_valid_vec3(pos)) {
+      throw new Error(`Invalid position at index ${idx}: expected 3 finite coordinates`)
+    }
+
+    const xyz = pos
+    const abc = cart_to_frac ? cart_to_frac(xyz) : ([0, 0, 0] as Vec3)
+
+    const force = force_data?.[idx]
+    const properties = is_valid_vec3(force) ? { force } : {}
+
+    return {
+      species: [{ element: elements[idx], occu: 1, oxidation_state: 0 }],
+      abc,
+      xyz,
+      label: `${elements[idx]}${idx + 1}`,
+      properties,
+    }
+  })
+
+  return lattice_matrix
+    ? {
+        sites,
+        lattice: {
+          matrix: lattice_matrix,
+          ...math.calc_lattice_params(lattice_matrix),
+          pbc: pbc ?? ([true, true, true] satisfies Pbc),
+        },
+      }
+    : { sites }
+}
+
+export const create_trajectory_frame = (
+  positions: number[][],
+  elements: ElementSymbol[],
+  lattice_matrix: math.Matrix3x3 | undefined,
+  pbc: Pbc | undefined,
+  step: number,
+  metadata: Record<string, unknown> = {},
+): TrajectoryFrame => ({
+  structure: create_structure(positions, elements, lattice_matrix, pbc),
+  step,
+  metadata,
+})
+
+// Shared utility to read ndarray data from binary format
+export const read_ndarray_from_view = (
+  view: DataView,
+  ref: { ndarray: unknown[] },
+): number[][] => {
+  const [shape, dtype, array_offset] = ref.ndarray as [number[], string, number]
+  const total = shape.reduce((product, dim_size) => product * dim_size, 1)
+  const data: number[] = []
+  let pos = array_offset
+
+  const readers = {
+    int64: {
+      bytes_per_element: 8,
+      read: () => {
+        const value = Number(view.getBigInt64(pos, true))
+        pos += 8
+        return value
+      },
+    },
+    int32: {
+      bytes_per_element: 4,
+      read: () => {
+        const value = view.getInt32(pos, true)
+        pos += 4
+        return value
+      },
+    },
+    float64: {
+      bytes_per_element: 8,
+      read: () => {
+        const value = view.getFloat64(pos, true)
+        pos += 8
+        return value
+      },
+    },
+    float32: {
+      bytes_per_element: 4,
+      read: () => {
+        const value = view.getFloat32(pos, true)
+        pos += 4
+        return value
+      },
+    },
+  }
+
+  const reader_config = readers[dtype as keyof typeof readers]
+  if (!reader_config) throw new Error(`Unsupported dtype: ${dtype}`)
+
+  if (!Number.isInteger(array_offset) || array_offset < 0) {
+    throw new Error(`Invalid array_offset: expected non-negative integer, got ${array_offset}`)
+  }
+
+  const bytes_needed = total * reader_config.bytes_per_element
+  if (array_offset + bytes_needed > view.byteLength) {
+    throw new Error(`Out-of-bounds read: array_offset + bytesNeeded exceeds view.byteLength`)
+  }
+
+  for (let idx = 0; idx < total; idx++) data.push(reader_config.read())
+
+  return shape.length === 1
+    ? [data]
+    : shape.length === 2
+      ? Array.from({ length: shape[0] }, (_, idx) =>
+          data.slice(idx * shape[1], (idx + 1) * shape[1]),
+        )
+      : (() => {
+          throw new Error(`Unsupported shape`)
+        })()
+}
+
+// Unified frame counting for XYZ
+export function count_xyz_frames(data: string): number {
+  if (!data || typeof data !== `string`) return 0
+  const lines = data.trim().split(/\r?\n/)
+  let frame_count = 0
+  let line_idx = 0
+
+  while (line_idx < lines.length) {
+    if (!lines[line_idx]?.trim()) {
+      line_idx++
+      continue
+    }
+
+    const num_atoms = parseInt(lines[line_idx].trim(), 10)
+    if (isNaN(num_atoms) || num_atoms <= 0 || line_idx + num_atoms + 2 > lines.length) {
+      line_idx++
+      continue
+    }
+
+    // Quick validation of first few atom lines
+    let valid_coords = 0
+    for (let idx = 0; idx < Math.min(num_atoms, 3); idx++) {
+      const parts = lines[line_idx + 2 + idx]?.trim().split(/\s+/)
+      if (parts?.length >= 4 && isNaN(parseInt(parts[0])) && parts[0].length <= 3) {
+        if (parts.slice(1, 4).every((coord) => !isNaN(parseFloat(coord)))) valid_coords++
+      }
+    }
+
+    if (valid_coords >= Math.min(num_atoms, 3)) {
+      frame_count++
+      line_idx += 2 + num_atoms
+    } else {
+      line_idx++
+    }
+  }
+
+  return frame_count
+}