Upload folder using huggingface_hub

LICENSE

@@ -0,0 +1,13 @@
+BSD 3-Clause License
+
+Copyright 2026 UChicago Argonne, LLC. All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS “AS IS” AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

README.md

@@ -0,0 +1,155 @@
+---
+license: bsd-3-clause
+language:
+- en
+tags:
+- pytorch
+- materials-science
+- crystallography
+- x-ray-diffraction
+- pxrd
+- convnext
+- arxiv:2603.23367
+datasets:
+- materials-project
+metrics:
+- accuracy
+- mae
+pipeline_tag: other
+---
+
+# AlphaDiffract — Open Weights
+
+[arXiv](https://arxiv.org/abs/2603.23367) | [GitHub](https://github.com/AdvancedPhotonSource/AlphaDiffract)
+
+**Automated crystallographic analysis of powder X-ray diffraction data.**
+
+AlphaDiffract is a multi-task 1D ConvNeXt model that takes a powder X-ray diffraction (PXRD) pattern and simultaneously predicts:
+
+| Output | Description |
+|---|---|
+| **Crystal system** | 7-class classification (Triclinic → Cubic) |
+| **Space group** | 230-class classification |
+| **Lattice parameters** | 6 values: a, b, c (Å), α, β, γ (°) |
+
+This release contains a **single model** trained exclusively on
+[Materials Project](https://next-gen.materialsproject.org/) structures
+(publicly available data). It is *not* the 10-model ensemble reported in
+the paper — see [Performance](#performance) for details.
+
+## Quick Start
+
+```bash
+pip install torch safetensors numpy
+```
+
+```python
+from model import AlphaDiffract
+import torch, numpy as np
+
+model = AlphaDiffract.from_pretrained(".", device="cpu")
+
+# 8192-point intensity pattern, normalized to [0, 100]
+pattern = np.load("my_pattern.npy").astype(np.float32)
+pattern = (pattern - pattern.min()) / (pattern.max() - pattern.min() + 1e-10) * 100.0
+x = torch.from_numpy(pattern).unsqueeze(0)
+
+with torch.no_grad():
+    out = model(x)
+
+cs_probs = torch.softmax(out["cs_logits"], dim=-1)
+sg_probs = torch.softmax(out["sg_logits"], dim=-1)
+lp = out["lp"]  # [a, b, c, alpha, beta, gamma]
+
+print("Crystal system:", AlphaDiffract.CRYSTAL_SYSTEMS[cs_probs.argmax().item()])
+print("Space group:   #", sg_probs.argmax().item() + 1)
+print("Lattice params:", lp[0].tolist())
+```
+
+See `example_inference.py` for a complete runnable example.
+
+## Files
+
+| File | Description |
+|---|---|
+| `model.safetensors` | Model weights (safetensors format, ~35 MB) |
+| `model.py` | Standalone model definition (pure PyTorch, no Lightning) |
+| `config.json` | Architecture and training hyperparameters |
+| `maxsub.json` | Space-group subgroup graph (230×230, used as a registered buffer) |
+| `example_inference.py` | End-to-end inference example |
+| `LICENSE` | BSD 3-Clause |
+
+
+## Input Format
+
+- **Length:** 8192 equally-spaced intensity values
+- **2θ range:** 5–20° (monochromatic, 20 keV)
+- **Preprocessing:** floor negatives at zero, then rescale to [0, 100]
+- **Shape:** `(batch, 8192)` or `(batch, 1, 8192)`
+
+## Architecture
+
+1D ConvNeXt backbone adapted from [Liu et al. (2022)](https://arxiv.org/abs/2201.03545):
+
+```
+Input (8192) → [ConvNeXt Block × 3 with AvgPool] → Flatten (560-d)
+  ├─ CS head:  MLP 560→2300→1150→7    (crystal system)
+  ├─ SG head:  MLP 560→2300→1150→230  (space group)
+  └─ LP head:  MLP 560→512→256→6      (lattice parameters, sigmoid-bounded)
+```
+
+- **Parameters:** 8,734,989
+- **Activation:** GELU
+- **Stochastic depth:** 0.3
+- **Head dropout:** 0.5
+
+## Performance
+
+This is a **single model** trained on Materials Project data only (no ICSD).
+Metrics on the best validation checkpoint (epoch 11):
+
+| Metric | Simulated Val | RRUFF (experimental) |
+|---|---|---|
+| Crystal system accuracy | 74.88% | 60.35% |
+| Space group accuracy | 57.31% | 38.28% |
+| Lattice parameter MAE | 2.71 | — |
+
+The paper reports higher numbers from a 10-model ensemble trained on
+Materials Project + ICSD combined data. This open-weights release covers
+only publicly available training data.
+
+## Training Details
+
+| | |
+|---|---|
+| **Data** | ~146k Materials Project structures, 100 GSAS-II simulations each |
+| **Augmentation** | Poisson + Gaussian noise, rescaled to [0, 100] |
+| **Optimizer** | AdamW (lr=2e-4, weight_decay=0.01) |
+| **Scheduler** | CyclicLR (triangular2, 6-epoch half-cycles) |
+| **Loss** | CE (crystal system) + CE + GEMD (space group) + MSE (lattice params) |
+| **Hardware** | 1× NVIDIA H100, float32 |
+| **Batch size** | 64 |
+
+## Citation
+
+```bibtex
+@article{andrejevic2026alphadiffract,
+  title   = {AlphaDiffract: Automated Crystallographic Analysis of Powder X-ray Diffraction Data},
+  author  = {Andrejevic, Nina and Du, Ming and Sharma, Hemant and Horwath, James P. and Luo, Aileen and Yin, Xiangyu and Prince, Michael and Toby, Brian H. and Cherukara, Mathew J.},
+  year    = {2026},
+  eprint  = {2603.23367},
+  archivePrefix = {arXiv},
+  primaryClass  = {cond-mat.mtrl-sci},
+  doi     = {10.48550/arXiv.2603.23367},
+  url     = {https://arxiv.org/abs/2603.23367}
+}
+```
+
+## License
+
+BSD 3-Clause — Copyright 2026 UChicago Argonne, LLC.
+
+## Links
+
+- [arXiv: 2603.23367](https://arxiv.org/abs/2603.23367)
+- [GitHub: OpenAlphaDiffract](https://github.com/AdvancedPhotonSource/AlphaDiffract)

config.json

@@ -0,0 +1,55 @@
+{
+  "model_type": "alphadiffract",
+  "architecture": "ConvNeXt1D-MultiTask",
+  "backbone": {
+    "dim_in": 8192,
+    "channels": [80, 80, 80],
+    "kernel_sizes": [100, 50, 25],
+    "strides": [5, 5, 5],
+    "dropout_rate": 0.3,
+    "ramped_dropout_rate": false,
+    "block_type": "convnext",
+    "pooling_type": "average",
+    "final_pool": true,
+    "use_batchnorm": false,
+    "activation": "gelu",
+    "output_type": "flatten",
+    "layer_scale_init_value": 0.0,
+    "drop_path_rate": 0.3
+  },
+  "heads": {
+    "head_dropout": 0.5,
+    "cs_hidden": [2300, 1150],
+    "sg_hidden": [2300, 1150],
+    "lp_hidden": [512, 256]
+  },
+  "tasks": {
+    "num_cs_classes": 7,
+    "num_sg_classes": 230,
+    "num_lp_outputs": 6,
+    "lp_bounds_min": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+    "lp_bounds_max": [500.0, 500.0, 500.0, 180.0, 180.0, 180.0],
+    "bound_lp_with_sigmoid": true
+  },
+  "training": {
+    "optimizer": "AdamW",
+    "lr": 0.0002,
+    "weight_decay": 0.01,
+    "scheduler": "CyclicLR",
+    "scheduler_mode": "triangular2",
+    "batch_size": 64,
+    "max_epochs": 100,
+    "precision": "float32",
+    "gemd_mu": 1.0,
+    "lambda_cs": 1.0,
+    "lambda_sg": 1.0,
+    "lambda_lp": 1.0
+  },
+  "preprocessing": {
+    "input_length": 8192,
+    "floor_at_zero": true,
+    "normalize_range": [0.0, 100.0],
+    "noise_poisson_range": [1.0, 100.0],
+    "noise_gaussian_range": [0.001, 0.1]
+  }
+}

example_inference.py

@@ -0,0 +1,44 @@
+"""
+Example: load AlphaDiffract and run inference on a PXRD pattern.
+
+Requirements:
+    pip install torch safetensors numpy
+"""
+
+import numpy as np
+import torch
+from model import AlphaDiffract
+
+# 1. Load model ---------------------------------------------------------------
+model = AlphaDiffract.from_pretrained(".", device="cpu")  # or "cuda"
+
+# 2. Prepare input -------------------------------------------------------------
+# The model expects an 8192-point PXRD intensity pattern normalized to [0, 100].
+# Replace this with your own data.
+pattern = np.random.rand(8192).astype(np.float32)  # placeholder
+
+# Normalize to [0, 100]
+pattern = (pattern - pattern.min()) / (pattern.max() - pattern.min() + 1e-10) * 100.0
+x = torch.from_numpy(pattern).unsqueeze(0)  # shape: (1, 8192)
+
+# 3. Inference -----------------------------------------------------------------
+with torch.no_grad():
+    out = model(x)
+
+cs_probs = torch.softmax(out["cs_logits"], dim=-1)
+sg_probs = torch.softmax(out["sg_logits"], dim=-1)
+lp = out["lp"]
+
+# 4. Results -------------------------------------------------------------------
+cs_idx = cs_probs.argmax(dim=-1).item()
+sg_idx = sg_probs.argmax(dim=-1).item()
+
+print(f"Crystal system : {AlphaDiffract.CRYSTAL_SYSTEMS[cs_idx]} "
+      f"({cs_probs[0, cs_idx]:.1%})")
+print(f"Space group    : #{sg_idx + 1} ({sg_probs[0, sg_idx]:.1%})")
+
+labels = ["a", "b", "c", "alpha", "beta", "gamma"]
+units = ["A", "A", "A", "deg", "deg", "deg"]
+print("Lattice params :")
+for name, val, unit in zip(labels, lp[0].tolist(), units):
+    print(f"  {name:>5s} = {val:8.3f} {unit}")

maxsub.json

@@ -0,0 +1,232 @@
+{
+    "1": [1],
+    "2": [1,2],
+    "3": [1,3,4,5],
+    "4": [1,4],
+    "5": [1,3,4,5],
+    "6": [1,6,7,8],
+    "7": [1,7,9],
+    "8": [1,6,7,8,9],
+    "9": [1,7,9],
+    "10": [2,3,6,10,11,12,13],
+    "11": [2,4,6,11,14],
+    "12": [2,5,8,10,11,12,13,14,15],
+    "13": [2,3,7,13,14,15],
+    "14": [2,4,7,14],
+    "15": [2,5,9,13,14,15],
+    "16": [3,16,17,21,22],
+    "17": [3,4,17,18,20],
+    "18": [3,4,18,19],
+    "19": [4,19],
+    "20": [4,5,17,18,19,20],
+    "21": [3,5,16,17,18,20,21,23,24],
+    "22": [5,20,21,22],
+    "23": [5,16,18,23],
+    "24": [5,17,19,24],
+    "25": [3,6,25,26,27,28,35,38,39,42],
+    "26": [4,6,7,26,29,31,36],
+    "27": [3,7,27,30,37],
+    "28": [3,6,7,28,29,30,31,32,40,41],
+    "29": [4,7,29,33],
+    "30": [3,7,30,34],
+    "31": [4,6,7,31,33],
+    "32": [3,7,32,33,34],
+    "33": [4,7,33],
+    "34": [3,7,34,43],
+    "35": [3,8,25,28,32,35,36,37,44,45,46],
+    "36": [4,8,9,26,29,31,33,36],
+    "37": [3,9,27,30,34,37],
+    "38": [5,6,8,25,26,30,31,38,40,44,46],
+    "39": [5,7,8,26,27,28,29,39,41,45,46],
+    "40": [5,6,9,28,31,33,34,40],
+    "41": [5,7,9,29,30,32,33,41],
+    "42": [5,8,35,36,37,38,39,40,41,42],
+    "43": [5,9,43],
+    "44": [5,8,25,31,34,44],
+    "45": [5,9,27,29,32,45],
+    "46": [5,8,9,26,28,30,33,46],
+    "47": [10,16,25,47,49,51,65,67,69],
+    "48": [13,16,34,48,70],
+    "49": [10,13,16,27,28,49,50,53,54,66,68],
+    "50": [13,16,30,32,48,50,52],
+    "51": [10,11,13,17,25,26,28,51,53,54,55,57,59,63,64],
+    "52": [13,14,17,30,33,34,52],
+    "53": [10,13,14,17,28,30,31,52,53,58,60],
+    "54": [13,14,17,27,29,32,52,54,56,60],
+    "55": [10,14,18,26,32,55,58,62],
+    "56": [13,14,18,27,33,56],
+    "57": [11,13,14,18,26,28,29,57,60,61,62],
+    "58": [10,14,18,31,34,58],
+    "59": [11,13,18,25,31,56,59,62],
+    "60": [13,14,18,29,30,33,60],
+    "61": [14,19,29,61],
+    "62": [11,14,19,26,31,33,62],
+    "63": [11,12,15,20,36,38,40,51,52,57,58,59,60,62,63],
+    "64": [12,14,15,20,36,39,41,53,54,55,56,57,60,61,62,64],
+    "65": [10,12,21,35,38,47,50,51,53,55,59,63,65,66,71,72,74],
+    "66": [10,15,21,37,40,48,49,52,53,56,58,66],
+    "67": [12,13,21,35,39,49,51,54,57,64,67,68,72,73,74],
+    "68": [13,15,21,37,41,50,52,54,60,68],
+    "69": [12,22,42,63,64,65,66,67,68,69],
+    "70": [15,22,43,70],
+    "71": [12,23,44,47,48,58,59,71],
+    "72": [12,15,23,45,46,49,50,55,56,57,60,72],
+    "73": [15,24,45,54,61,73],
+    "74": [12,15,24,44,46,51,52,53,62,74],
+    "75": [3,75,77,79],
+    "76": [4,76,78],
+    "77": [3,76,77,78,80],
+    "78": [4,76,78],
+    "79": [5,75,77,79],
+    "80": [5,76,78,80],
+    "81": [3,81,82],
+    "82": [5,81,82],
+    "83": [10,75,81,83,84,85,87],
+    "84": [10,77,81,84,86],
+    "85": [13,75,81,85,86],
+    "86": [13,77,81,86,88],
+    "87": [12,79,82,83,84,85,86,87],
+    "88": [15,80,82,88],
+    "89": [16,21,75,89,90,93,97],
+    "90": [18,21,75,90,94],
+    "91": [17,20,76,91,92,95],
+    "92": [19,20,76,92,96],
+    "93": [16,21,77,91,93,94,95,98],
+    "94": [18,21,77,92,94,96],
+    "95": [17,20,78,91,95,96],
+    "96": [19,20,78,92,96],
+    "97": [22,23,79,89,90,93,94,97],
+    "98": [22,24,80,91,92,95,96,98],
+    "99": [25,35,75,99,100,101,103,105,107,108],
+    "100": [32,35,75,100,102,104,106],
+    "101": [27,35,77,101,105,106],
+    "102": [34,35,77,102,109,110],
+    "103": [27,37,75,103,104],
+    "104": [34,37,75,104],
+    "105": [25,37,77,101,102,105],
+    "106": [32,37,77,106],
+    "107": [42,44,79,99,102,104,105,107],
+    "108": [42,45,79,100,101,103,106,108],
+    "109": [43,44,80,109],
+    "110": [43,45,80,110],
+    "111": [16,35,81,111,112,115,117,119,120],
+    "112": [16,37,81,112,116,118],
+    "113": [18,35,81,113,114],
+    "114": [18,37,81,114],
+    "115": [21,25,81,111,113,115,116,121],
+    "116": [21,27,81,112,114,116],
+    "117": [21,32,81,117,118],
+    "118": [21,34,81,118,122],
+    "119": [22,44,82,115,118,119],
+    "120": [22,45,82,116,117,120],
+    "121": [23,42,82,111,112,113,114,121],
+    "122": [24,43,82,122],
+    "123": [47,65,83,89,99,111,115,123,124,125,127,129,131,132,139,140],
+    "124": [49,66,83,89,103,112,116,124,126,128,130],
+    "125": [50,67,85,89,100,111,117,125,126,133,134],
+    "126": [48,68,85,89,104,112,118,126],
+    "127": [55,65,83,90,100,113,117,127,128,135,136],
+    "128": [58,66,83,90,104,114,118,128],
+    "129": [59,67,85,90,99,113,115,129,130,137,138],
+    "130": [56,68,85,90,103,114,116,130],
+    "131": [47,66,84,93,105,112,115,131,132,134,136,138],
+    "132": [49,65,84,93,101,111,116,131,132,133,135,137],
+    "133": [50,68,86,93,106,112,117,133],
+    "134": [48,67,86,93,102,111,118,134,141,142],
+    "135": [55,66,84,94,106,114,117,135],
+    "136": [58,65,84,94,102,113,118,136],
+    "137": [59,68,86,94,105,114,115,137],
+    "138": [56,67,86,94,101,113,116,138],
+    "139": [69,71,87,97,107,119,121,123,126,128,129,131,134,136,137,139],
+    "140": [69,72,87,97,108,120,121,124,125,127,130,132,133,135,138,140],
+    "141": [70,74,88,98,109,119,122,141],
+    "142": [70,73,88,98,110,120,122,142],
+    "143": [1,143,144,145,146],
+    "144": [1,144,145],
+    "145": [1,144,145],
+    "146": [1,143,144,145,146],
+    "147": [2,143,147,148],
+    "148": [2,146,147,148],
+    "149": [5,143,149,150,151,153,155],
+    "150": [5,143,149,150,152,154],
+    "151": [5,144,151,152,153],
+    "152": [5,144,151,152,154],
+    "153": [5,145,151,153,154],
+    "154": [5,145,152,153,154],
+    "155": [5,146,150,152,154,155],
+    "156": [8,143,156,157,158],
+    "157": [8,143,156,157,159,160],
+    "158": [9,143,158,159],
+    "159": [9,143,158,159,161],
+    "160": [8,146,156,160,161],
+    "161": [9,146,158,161],
+    "162": [12,147,149,157,162,163,164,166],
+    "163": [15,147,149,159,163,165,167],
+    "164": [12,147,150,156,162,164,165],
+    "165": [15,147,150,158,163,165],
+    "166": [12,148,155,160,164,166,167],
+    "167": [15,148,155,161,165,167],
+    "168": [3,143,168,171,172,173],
+    "169": [4,144,169,170],
+    "170": [4,145,169,170],
+    "171": [3,145,169,171,172],
+    "172": [3,144,170,171,172],
+    "173": [4,143,169,170,173],
+    "174": [6,143,174],
+    "175": [10,147,168,174,175,176],
+    "176": [11,147,173,174,176],
+    "177": [21,149,150,168,177,180,181,182],
+    "178": [20,151,152,169,178,179],
+    "179": [20,153,154,170,178,179],
+    "180": [21,153,154,171,178,180,181],
+    "181": [21,151,152,172,179,180,181],
+    "182": [20,149,150,173,178,179,182],
+    "183": [35,156,157,168,183,184,185,186],
+    "184": [37,158,159,168,184],
+    "185": [36,157,158,173,185,186],
+    "186": [36,156,159,173,185,186],
+    "187": [38,149,156,174,187,188,189],
+    "188": [40,149,158,174,188,190],
+    "189": [38,150,157,174,187,189,190],
+    "190": [40,150,159,174,188,190],
+    "191": [65,162,164,175,177,183,187,189,191,192,193,194],
+    "192": [66,163,165,175,177,184,188,190,192],
+    "193": [63,162,165,176,182,185,188,189,193,194],
+    "194": [63,163,164,176,182,186,187,190,193,194],
+    "195": [16,146,196,197,199],
+    "196": [22,146,195,198],
+    "197": [23,146,195],
+    "198": [19,146],
+    "199": [24,146,198],
+    "200": [47,148,195,202,204,206],
+    "201": [48,148,195,203],
+    "202": [69,148,196,200,201,205],
+    "203": [70,148,196],
+    "204": [71,148,197,200,201],
+    "205": [61,148,198],
+    "206": [73,148,199,205],
+    "207": [89,155,195,209,211],
+    "208": [93,155,195,210,214],
+    "209": [97,155,196,207,208],
+    "210": [98,155,196,212,213],
+    "211": [97,155,197,207,208],
+    "212": [96,155,198],
+    "213": [92,155,198],
+    "214": [98,155,199,212,213],
+    "215": [111,160,195,216,217,219],
+    "216": [119,160,196,215],
+    "217": [121,160,197,215,218],
+    "218": [112,161,195,220],
+    "219": [120,161,196,218],
+    "220": [122,161,199],
+    "221": [123,166,200,207,215,225,226,229],
+    "222": [126,167,201,207,218],
+    "223": [131,167,200,208,218,230],
+    "224": [134,166,201,208,215,227,228],
+    "225": [139,166,202,209,216,221,224],
+    "226": [140,167,202,209,219,222,223],
+    "227": [141,166,203,210,216],
+    "228": [142,167,203,210,219],
+    "229": [139,166,204,211,217,221,222,223,224],
+    "230": [142,167,206,214,220]
+}

model.py

@@ -0,0 +1,408 @@
+"""
+This file is self-contained: download it alongside `model.safetensors`,
+`config.json`, and `maxsub.json` to load and run the model.
+"""
+
+import json
+from collections import deque
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ---------------------------------------------------------------------------
+# Utility: DropPath (Stochastic Depth)
+# ---------------------------------------------------------------------------
+def drop_path(
+    x: torch.Tensor, drop_prob: float = 0.0, training: bool = False
+) -> torch.Tensor:
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor = random_tensor.floor()
+    return x.div(keep_prob) * random_tensor
+
+
+class DropPath(nn.Module):
+    def __init__(self, drop_prob: float = 0.0):
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return drop_path(x, self.drop_prob, self.training)
+
+
+# ---------------------------------------------------------------------------
+# ConvNeXt 1D Block
+# ---------------------------------------------------------------------------
+class ConvNeXtBlock1D(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        kernel_size: int,
+        drop_path: float,
+        layer_scale_init_value: float,
+        activation: nn.Module,
+    ):
+        super().__init__()
+        self.dwconv = nn.Conv1d(
+            dim, dim, kernel_size=kernel_size, padding="same", groups=dim
+        )
+        self.pwconv1 = nn.Linear(dim, 4 * dim)
+        self.act = activation() if isinstance(activation, type) else activation
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.gamma = (
+            nn.Parameter(layer_scale_init_value * torch.ones(dim))
+            if layer_scale_init_value > 0
+            else None
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shortcut = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 1)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = x * self.gamma
+        x = x.permute(0, 2, 1)
+        x = shortcut + self.drop_path(x)
+        return x
+
+
+class ConvNextBlock1DAdaptor(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int,
+        dropout: float,
+        use_batchnorm: bool,
+        activation: nn.Module,
+        layer_scale_init_value: float,
+        drop_path_rate: float,
+        block_type: str,
+    ):
+        super().__init__()
+        if in_channels != out_channels:
+            act = activation() if isinstance(activation, type) else activation
+            self.pwconv = nn.Sequential(nn.Linear(in_channels, out_channels), act)
+        else:
+            self.pwconv = None
+
+        if block_type == "convnext":
+            self.block = ConvNeXtBlock1D(
+                dim=out_channels,
+                kernel_size=kernel_size,
+                drop_path=drop_path_rate,
+                layer_scale_init_value=layer_scale_init_value,
+                activation=activation,
+            )
+        else:
+            self.block = None
+
+        if stride > 1:
+            self.reduction_pool = nn.AvgPool1d(kernel_size=stride, stride=stride)
+        else:
+            self.reduction_pool = None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.pwconv is not None:
+            x = x.permute(0, 2, 1)
+            x = self.pwconv(x)
+            x = x.permute(0, 2, 1)
+        if self.block is not None:
+            x = self.block(x)
+        if self.reduction_pool is not None:
+            x = self.reduction_pool(x)
+        return x
+
+
+# ---------------------------------------------------------------------------
+# MLP head builder
+# ---------------------------------------------------------------------------
+def make_mlp(
+    input_dim: int,
+    hidden_dims: Optional[Tuple[int, ...]],
+    output_dim: int,
+    dropout: float = 0.2,
+    output_activation: Optional[nn.Module] = None,
+) -> nn.Module:
+    layers: List[nn.Module] = []
+    last = input_dim
+    if hidden_dims is not None and len(hidden_dims) > 0:
+        for hd in hidden_dims:
+            layers.extend([nn.Linear(last, hd), nn.ReLU()])
+            if dropout and dropout > 0:
+                layers.append(nn.Dropout(dropout))
+            last = hd
+    layers.append(nn.Linear(last, output_dim))
+    if output_activation is not None:
+        layers.append(output_activation)
+    return nn.Sequential(*layers)
+
+
+# ---------------------------------------------------------------------------
+# Backbone
+# ---------------------------------------------------------------------------
+class MultiscaleCNNBackbone1D(nn.Module):
+    def __init__(
+        self,
+        dim_in: int,
+        channels: Tuple[int, ...],
+        kernel_sizes: Tuple[int, ...],
+        strides: Tuple[int, ...],
+        dropout_rate: float,
+        ramped_dropout_rate: bool,
+        block_type: str,
+        pooling_type: str,
+        final_pool: bool,
+        use_batchnorm: bool,
+        activation: nn.Module,
+        output_type: str,
+        layer_scale_init_value: float,
+        drop_path_rate: float,
+    ):
+        super().__init__()
+        assert len(channels) == len(kernel_sizes) == len(strides)
+        self.dim_in = dim_in
+        self.output_type = output_type
+
+        if ramped_dropout_rate:
+            dropout_per_stage = torch.linspace(
+                0.0, dropout_rate, steps=len(channels)
+            ).tolist()
+        else:
+            dropout_per_stage = [dropout_rate] * len(channels)
+
+        if pooling_type == "average":
+            pool_cls = nn.AvgPool1d
+            pool_kwargs = {"kernel_size": 3, "stride": 2}
+        elif pooling_type == "max":
+            pool_cls = nn.MaxPool1d
+            pool_kwargs = {"kernel_size": 2, "stride": 2}
+        else:
+            raise ValueError(f"Invalid pooling_type '{pooling_type}'")
+
+        layers: List[nn.Module] = []
+        in_ch = 1
+        for i, (out_ch, k, s) in enumerate(zip(channels, kernel_sizes, strides)):
+            stage_block = ConvNextBlock1DAdaptor(
+                in_channels=in_ch,
+                out_channels=out_ch,
+                kernel_size=k,
+                stride=s,
+                dropout=dropout_per_stage[i],
+                use_batchnorm=use_batchnorm,
+                activation=activation,
+                layer_scale_init_value=layer_scale_init_value,
+                drop_path_rate=drop_path_rate,
+                block_type=block_type,
+            )
+            layers.append(stage_block)
+            if i < len(channels) - 1 or final_pool:
+                layers.append(pool_cls(**pool_kwargs))
+            in_ch = out_ch
+
+        self.net = nn.Sequential(*layers)
+
+        if self.output_type == "gap":
+            self.dim_output = channels[-1]
+        elif self.output_type == "flatten":
+            with torch.no_grad():
+                dummy = torch.zeros(1, 1, self.dim_in)
+                out = self.net(dummy)
+                self.dim_output = int(out.shape[1] * out.shape[2])
+        else:
+            raise ValueError(f"Invalid output_type '{self.output_type}'")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.ndim == 2:
+            x = x[:, None, :]
+        x = self.net(x)
+        if self.output_type == "gap":
+            x = x.mean(dim=-1)
+        else:
+            x = x.reshape(x.shape[0], -1)
+        return x
+
+
+# ---------------------------------------------------------------------------
+# GEMD distance-matrix utilities
+# ---------------------------------------------------------------------------
+def _build_distance_matrix_from_maxsub_lut(
+    maxsub_lut: Dict[str, List[int]],
+    num_sg_classes: int,
+) -> torch.Tensor:
+    adjacency: List[set] = [set() for _ in range(num_sg_classes)]
+    for key, neighbors in maxsub_lut.items():
+        src = int(key) - 1
+        for raw_dst in neighbors:
+            dst = int(raw_dst) - 1
+            adjacency[src].add(dst)
+            adjacency[dst].add(src)
+
+    distance_matrix = torch.zeros(
+        (num_sg_classes, num_sg_classes), dtype=torch.float32
+    )
+    for src in range(num_sg_classes):
+        dists = [-1] * num_sg_classes
+        dists[src] = 0
+        queue = deque([src])
+        while queue:
+            cur = queue.popleft()
+            for nxt in adjacency[cur]:
+                if dists[nxt] == -1:
+                    dists[nxt] = dists[cur] + 1
+                    queue.append(nxt)
+        distance_matrix[src] = torch.tensor(dists, dtype=torch.float32)
+    return distance_matrix
+
+
+def load_gemd_distance_matrix(
+    path: str, num_sg_classes: int = 230
+) -> torch.Tensor:
+    with open(path, "r", encoding="utf-8") as f:
+        payload: Any = json.load(f)
+    if isinstance(payload, dict) and all(str(k).isdigit() for k in payload.keys()):
+        return _build_distance_matrix_from_maxsub_lut(payload, num_sg_classes)
+    elif isinstance(payload, list):
+        return torch.as_tensor(payload, dtype=torch.float32)
+    raise ValueError(f"Could not parse GEMD data from {path}")
+
+
+# ---------------------------------------------------------------------------
+# Full model
+# ---------------------------------------------------------------------------
+class AlphaDiffract(nn.Module):
+    """
+    AlphaDiffract: multi-task 1D ConvNeXt for powder X-ray diffraction
+    pattern analysis.
+
+    Predicts crystal system (7 classes), space group (230 classes), and
+    lattice parameters (6 values: a, b, c, alpha, beta, gamma).
+    """
+
+    CRYSTAL_SYSTEMS = [
+        "Triclinic",
+        "Monoclinic",
+        "Orthorhombic",
+        "Tetragonal",
+        "Trigonal",
+        "Hexagonal",
+        "Cubic",
+    ]
+
+    def __init__(self, config: dict, maxsub_path: Optional[str] = None):
+        super().__init__()
+        bb = config["backbone"]
+        heads = config["heads"]
+        tasks = config["tasks"]
+
+        activation = nn.GELU
+
+        self.backbone = MultiscaleCNNBackbone1D(
+            dim_in=bb["dim_in"],
+            channels=tuple(bb["channels"]),
+            kernel_sizes=tuple(bb["kernel_sizes"]),
+            strides=tuple(bb["strides"]),
+            dropout_rate=bb["dropout_rate"],
+            ramped_dropout_rate=bb["ramped_dropout_rate"],
+            block_type=bb["block_type"],
+            pooling_type=bb["pooling_type"],
+            final_pool=bb["final_pool"],
+            use_batchnorm=bb["use_batchnorm"],
+            activation=activation,
+            output_type=bb["output_type"],
+            layer_scale_init_value=bb["layer_scale_init_value"],
+            drop_path_rate=bb["drop_path_rate"],
+        )
+        feat_dim = self.backbone.dim_output
+
+        self.cs_head = make_mlp(
+            feat_dim, tuple(heads["cs_hidden"]), tasks["num_cs_classes"],
+            dropout=heads["head_dropout"],
+        )
+        self.sg_head = make_mlp(
+            feat_dim, tuple(heads["sg_hidden"]), tasks["num_sg_classes"],
+            dropout=heads["head_dropout"],
+        )
+        self.lp_head = make_mlp(
+            feat_dim, tuple(heads["lp_hidden"]), tasks["num_lp_outputs"],
+            dropout=heads["head_dropout"],
+        )
+
+        self.bound_lp_with_sigmoid = tasks["bound_lp_with_sigmoid"]
+        self.register_buffer(
+            "lp_min",
+            torch.tensor(tasks["lp_bounds_min"], dtype=torch.float32),
+        )
+        self.register_buffer(
+            "lp_max",
+            torch.tensor(tasks["lp_bounds_max"], dtype=torch.float32),
+        )
+
+        if maxsub_path is not None:
+            gemd = load_gemd_distance_matrix(maxsub_path)
+            self.register_buffer("gemd_distance_matrix", gemd)
+        else:
+            self.gemd_distance_matrix = None
+
+    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
+        """
+        Args:
+            x: PXRD pattern tensor of shape ``(batch, 8192)`` or
+               ``(batch, 1, 8192)``, intensity-normalized to [0, 100].
+
+        Returns:
+            Dict with keys ``cs_logits``, ``sg_logits``, ``lp``.
+        """
+        feats = self.backbone(x)
+        cs_logits = self.cs_head(feats)
+        sg_logits = self.sg_head(feats)
+        lp = self.lp_head(feats)
+        if self.bound_lp_with_sigmoid:
+            lp = torch.sigmoid(lp) * (self.lp_max - self.lp_min) + self.lp_min
+        return {"cs_logits": cs_logits, "sg_logits": sg_logits, "lp": lp}
+
+    # -- convenience loaders ------------------------------------------------
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_dir: str,
+        device: str = "cpu",
+    ) -> "AlphaDiffract":
+        """Load model from a directory containing config.json,
+        model.safetensors, and maxsub.json."""
+        model_dir = Path(model_dir)
+        with open(model_dir / "config.json", "r") as f:
+            config = json.load(f)
+
+        maxsub_path = model_dir / "maxsub.json"
+        model = cls(
+            config,
+            maxsub_path=str(maxsub_path) if maxsub_path.exists() else None,
+        )
+
+        weights_path = model_dir / "model.safetensors"
+        if weights_path.exists():
+            from safetensors.torch import load_file
+            state_dict = load_file(str(weights_path), device=device)
+        else:
+            # Fallback to PyTorch format
+            pt_path = model_dir / "model.pt"
+            state_dict = torch.load(str(pt_path), map_location=device, weights_only=True)
+
+        model.load_state_dict(state_dict)
+        model.to(device)
+        model.eval()
+        return model

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d81f2716ef8c78683d52ac51afff3eaf160c0b2b410685d0c90299bc2fd58ed
+size 35155332