java-topology/defects/pytorch-geometric/unit/PyGSmilesIndexAlgorithm.java

262 lines
11 KiB
Java
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

package unit;
import java.util.*;
/**
* pytorch-geometric-0001: smiles.py list.index() O(L×A) → O(A) with dict
*
* Simulates the defective pattern from torch_geometric/utils/smiles.py:
* from_rdmol() calls x_map[key].index(value) for each atom feature,
* where x_map[key] is a List. This is O(L) per call, called A×9 times
* per molecule, M molecules total → O(M × A × L).
*
* Fixed: pre-build x_idx[key] = {value → index} HashMap → O(1) lookup.
*
* Run: javac -d . PyGSmilesIndexAlgorithm.java && java unit.PyGSmilesIndexAlgorithm
*/
public class PyGSmilesIndexAlgorithm {
// ---- Result ----
static class Result {
final int slowOps;
final int fastOps;
Result(int slowOps, int fastOps) {
this.slowOps = slowOps;
this.fastOps = fastOps;
}
}
// ---- Defective: list.index() O(L) per lookup ----
static class DefectiveMolProcessor {
// Simulates x_map['atomic_num'] = list(range(0, 119))
final List<Integer> atomicNumList;
// Simulates x_map['chirality'] = 9 entries
final List<String> chiralityList;
// Simulates e_map['bond_type'] = 22 entries
final List<String> bondTypeList;
int opCount = 0;
DefectiveMolProcessor() {
atomicNumList = new ArrayList<>();
for (int i = 0; i < 119; i++) atomicNumList.add(i);
chiralityList = new ArrayList<>();
String[] chiralities = {"CHI_UNSPECIFIED","CHI_TETRAHEDRAL_CW","CHI_TETRAHEDRAL_CCW",
"CHI_OTHER","CHI_TETRAHEDRAL","CHI_ALLENE","CHI_SQUAREPLANAR",
"CHI_TRIGONALBIPYRAMIDAL","CHI_OCTAHEDRAL"};
chiralityList.addAll(Arrays.asList(chiralities));
bondTypeList = new ArrayList<>();
for (int i = 0; i < 22; i++) bondTypeList.add("BOND_TYPE_" + i);
}
// O(L) list scan — mirrors list.index()
int listIndex(List<?> list, Object value) {
for (int i = 0; i < list.size(); i++) {
opCount++;
if (list.get(i).equals(value)) return i;
}
throw new NoSuchElementException("Value not found: " + value);
}
// Process one molecule: A atoms, B bonds
// Each atom: 3 list.index() calls (atomic_num, chirality, dummy)
// Each bond: 1 list.index() call (bond_type)
void processMolecule(int[] atomNums, String[] chiralTags, String[] bondTypes) {
for (int i = 0; i < atomNums.length; i++) {
listIndex(atomicNumList, atomNums[i]);
listIndex(chiralityList, chiralTags[i % chiralTags.length]);
// Simulate remaining 7 atom feature lookups (simplified):
// degree, formal_charge, num_hs, num_radical, hybridization,
// is_aromatic, is_in_ring — all small lists, count as 7 ops each
opCount += 7;
}
for (String bt : bondTypes) {
listIndex(bondTypeList, bt);
}
}
Result run(int[][] atomNumsPerMol, String[] chiralTags, String[] bondTypes) {
opCount = 0;
for (int[] atomNums : atomNumsPerMol) {
processMolecule(atomNums, chiralTags, bondTypes);
}
return new Result(opCount, -1);
}
}
// ---- Fixed: HashMap O(1) lookup ----
static class FixedMolProcessor {
final Map<Integer, Integer> atomicNumIdx;
final Map<String, Integer> chiralityIdx;
final Map<String, Integer> bondTypeIdx;
int opCount = 0;
FixedMolProcessor() {
// Build index maps once at module load
atomicNumIdx = new HashMap<>();
for (int i = 0; i < 119; i++) atomicNumIdx.put(i, i);
chiralityIdx = new HashMap<>();
String[] chiralities = {"CHI_UNSPECIFIED","CHI_TETRAHEDRAL_CW","CHI_TETRAHEDRAL_CCW",
"CHI_OTHER","CHI_TETRAHEDRAL","CHI_ALLENE","CHI_SQUAREPLANAR",
"CHI_TRIGONALBIPYRAMIDAL","CHI_OCTAHEDRAL"};
for (int i = 0; i < chiralities.length; i++) chiralityIdx.put(chiralities[i], i);
bondTypeIdx = new HashMap<>();
for (int i = 0; i < 22; i++) bondTypeIdx.put("BOND_TYPE_" + i, i);
}
int mapGet(Map<?, Integer> map, Object key) {
opCount++; // one op per lookup (hash + compare)
return map.get(key);
}
void processMolecule(int[] atomNums, String[] chiralTags, String[] bondTypes) {
for (int i = 0; i < atomNums.length; i++) {
mapGet(atomicNumIdx, atomNums[i]);
mapGet(chiralityIdx, chiralTags[i % chiralTags.length]);
opCount += 7; // 7 other O(1) dict lookups
}
for (String bt : bondTypes) {
mapGet(bondTypeIdx, bt);
}
}
Result run(int[][] atomNumsPerMol, String[] chiralTags, String[] bondTypes) {
opCount = 0;
for (int[] atomNums : atomNumsPerMol) {
processMolecule(atomNums, chiralTags, bondTypes);
}
return new Result(-1, opCount);
}
}
// ---- check helper ----
static int passed = 0;
static int total = 0;
static void check(String name, boolean condition) {
total++;
if (condition) {
passed++;
System.out.println(" PASS: " + name);
} else {
System.out.println(" FAIL: " + name);
}
}
public static void main(String[] args) {
System.out.println("pytorch-geometric-0001: smiles.py list.index() O(L*A) -> O(A)");
System.out.println();
// Build test data: simulate molecules with varying atom counts
// Atoms use atomic nums spread across the list to trigger long scans
String[] chiralTags = {"CHI_UNSPECIFIED", "CHI_TETRAHEDRAL_CW", "CHI_OTHER"};
String[] bondTypes = {"BOND_TYPE_0", "BOND_TYPE_10", "BOND_TYPE_21"};
// ---- Test 1: Small molecule (10 atoms, 9 bonds) ----
{
int M = 1;
int[][] atomNums = new int[M][];
// Use high atomic numbers to stress the list scan
atomNums[0] = new int[]{100, 110, 90, 50, 118, 60, 80, 30, 5, 1};
DefectiveMolProcessor slow = new DefectiveMolProcessor();
Result slowR = slow.run(atomNums, chiralTags, bondTypes);
FixedMolProcessor fast = new FixedMolProcessor();
Result fastR = fast.run(atomNums, chiralTags, bondTypes);
System.out.println("Test 1: 1 molecule, 10 atoms");
System.out.println(" Slow ops: " + slowR.slowOps);
System.out.println(" Fast ops: " + fastR.fastOps);
check("slow > fast (list scan > dict lookup)", slowR.slowOps > fastR.fastOps);
check("slow ops > 300 (O(L) scans dominate)", slowR.slowOps > 300);
}
// ---- Test 2: Large dataset (1000 molecules, 18 atoms each) ----
{
int M = 1000;
int A = 18;
int[][] atomNums = new int[M][A];
Random rng = new Random(42);
for (int m = 0; m < M; m++) {
for (int a = 0; a < A; a++) {
atomNums[m][a] = rng.nextInt(119);
}
}
DefectiveMolProcessor slow = new DefectiveMolProcessor();
Result slowR = slow.run(atomNums, chiralTags, bondTypes);
FixedMolProcessor fast = new FixedMolProcessor();
Result fastR = fast.run(atomNums, chiralTags, bondTypes);
System.out.println("Test 2: 1000 molecules, 18 atoms each");
System.out.println(" Slow ops: " + slowR.slowOps);
System.out.println(" Fast ops: " + fastR.fastOps);
double ratio = (double) slowR.slowOps / fastR.fastOps;
System.out.printf(" Ratio slow/fast: %.1fx%n", ratio);
check("slow > 5x fast ops (O(L) overhead confirmed)", ratio > 5.0);
check("slow ops scale with L=119 (>500k for M=1000,A=18)", slowR.slowOps > 500_000);
}
// ---- Test 3: Correctness — same results ----
{
int[][] singleMol = {{6, 7, 8, 1, 16}}; // C, N, O, H, S
DefectiveMolProcessor slow = new DefectiveMolProcessor();
FixedMolProcessor fast = new FixedMolProcessor();
// Verify same index returned for known values
int slowIdx = slow.listIndex(slow.atomicNumList, 6); // Carbon = index 6
int fastIdx = fast.atomicNumIdx.get(6); // Carbon = index 6
System.out.println("Test 3: correctness check for atomic_num=6 (Carbon)");
System.out.println(" Slow index: " + slowIdx + ", Fast index: " + fastIdx);
check("slow and fast return same index (6)", slowIdx == fastIdx && fastIdx == 6);
int slowChiral = slow.listIndex(slow.chiralityList, "CHI_TETRAHEDRAL_CW");
int fastChiral = fast.chiralityIdx.get("CHI_TETRAHEDRAL_CW");
System.out.println("Test 4: correctness check for chirality=CHI_TETRAHEDRAL_CW");
check("slow and fast return same chirality index (1)", slowChiral == fastChiral && fastChiral == 1);
}
// ---- Test 4: O(n^2) scaling — double molecule count doubles slow more ----
{
int A = 10;
String[] chiralTagsT = {"CHI_UNSPECIFIED"};
String[] bondTypesT = {"BOND_TYPE_21"}; // index 21 = worst case scan
int M1 = 100;
int M2 = 200;
int[][] atomNums1 = new int[M1][A];
int[][] atomNums2 = new int[M2][A];
for (int m = 0; m < M2; m++) {
for (int a = 0; a < A; a++) {
int v = 100 + (m % 19); // high atomic nums to stress scan
if (m < M1) atomNums1[m][a] = v;
atomNums2[m][a] = v;
}
}
DefectiveMolProcessor slow1 = new DefectiveMolProcessor();
DefectiveMolProcessor slow2 = new DefectiveMolProcessor();
slow1.run(atomNums1, chiralTagsT, bondTypesT);
slow2.run(atomNums2, chiralTagsT, bondTypesT);
System.out.println("Test 5: O(n) scaling — 2x molecules → ~2x ops");
System.out.println(" M=" + M1 + " ops: " + slow1.opCount);
System.out.println(" M=" + M2 + " ops: " + slow2.opCount);
double scalingRatio = (double) slow2.opCount / slow1.opCount;
System.out.printf(" Scaling: %.2fx (should be ~2.0)%n", scalingRatio);
check("slow ops scale ~2x when molecules doubled", scalingRatio > 1.9 && scalingRatio < 2.1);
}
System.out.println();
System.out.println(passed + "/" + total + " PASS");
}
}