Update readme

Author: Linked-Liszt

README.md

@@ -1,30 +1,36 @@
 # OpenAlphaDiffract
 
-OpenAlphaDiffract is an open-source implementation of the AlphaDiffract research project. It provides a reproducible pipeline to:
+OpenAlphaDiffract is an open-source implementation of the [AlphaDiffract](https://arxiv.org/abs/2603.23367) research project. It provides a reproducible pipeline to:
 - Create a dataset from the Materials Project
 - Simulate powder diffraction patterns from those structures
 - Train and evaluate models on the generated dataset
 - Run an inference web app to try out the model
 
 
-## Dataset Pipeline Overview
+## Pretrained Model & Demo
 
-1. Acquire CIFs (downloader.Dockerfile)
+- **Model weights:** [linked-liszt/OpenAlphaDiffract](https://huggingface.co/linked-liszt/OpenAlphaDiffract) on Hugging Face
+- **Live demo:** [linked-liszt/OpenAlphaDiffract-UI](https://huggingface.co/spaces/linked-liszt/OpenAlphaDiffract-UI) on Hugging Face Spaces
+
+
+## Dataset Pipeline 
+
+1. Acquire CIFs (`docker/downloader.Dockerfile`)
     - Uses the Materials Project API to fetch crystal structures as CIF files
     - Configurable via `configs/download.yaml`
     - Filters structures by checking conventional cell consistency across multiple angle tolerances. This filters ~4.4% of MP structures as of 10/22/2025. 
 
-2. GSAS-II XRD Simulation (simulator.Dockerfile)
+2. GSAS-II XRD Simulation (`docker/simulator.Dockerfile`)
     - Generates synthetic powder diffraction patterns from CIFs
     - Configurable via `configs/simulator.yaml` (e.g., instrument file, noise ranges, job parallelism)
     - Creates .npy files with simulated pattern and metadata ready to be consumed by the training system
 
-3. Open Alpha Diffract Training (trainer.Dockerfile)
+3. Open Alpha Diffract Training (`docker/trainer.Dockerfile`)
    - Trains the multi-task AlphaDiffract model on the generated dataset
    - Configurable via `configs/trainer.docker.yaml` or `configs/trainer.local.yaml`
    - Logs checkpoints and metrics (CSV and optional MLflow)
 
-4. XRD Inference Web App (ui.Dockerfile)
+4. XRD Inference Web App (`docker/ui.Dockerfile`)
    - FastAPI service for model inference with a React frontend
    - Accepts processed XRD patterns and returns predictions via `/api/predict`
    - Serves the built frontend from the same container
@@ -39,8 +45,8 @@ Prerequisites:
 > [!WARNING]
 > Building the dataset and training will take a significant amount of space and computational resources: 
 > - Expect to use around 1TB+ of space in total to replicate the paper's 100-variation dataset
-> - We recommend running simulation with ~100 processes in parallel. For reference [XYZ] this should take [XYZ hours]. 
-> - Training took [XYZ hours] on [XYZ hardware]
+> - We recommend running simulation with ~100+ processes in parallel. For reference, simulation took ~18 hours on 2x AMD EPYC 7742 (128 processes). 
+> - Training took ~15 hours on a single H100 GPU.
 
 
 Setup:
@@ -50,11 +56,11 @@ Setup:
    - Optionally set `UID` and `GID` so the containers write files as your user. 
 
 2. Download CIFs:
-   - `scripts/download.sh` (or docker compose run --rm trainer)
+   - `scripts/download.sh` (or `docker compose run --rm downloader`)
    - CIFs will be written to `./data/raw_cif`
 
 3. Simulate diffraction patterns:
-   - `scripts/simulate.sh` (or docker compose run --rm simulator)
+   - `scripts/simulate.sh` (or `docker compose run --rm simulator`)
    - Patterns will be written to `./data/dataset`
    - Errors (if any) go to `./data/error_logs`
 
@@ -65,6 +71,7 @@ Setup:
 5. Run the inference UI:
    - Move a model checkpoint to `./src/ui/models/xrd_model.ckpt`
    - `docker compose up ui`
+   - Open `http://localhost:7860`
 
 Notes:
 - You can pass extra CLI args to the simulator via `scripts/simulate.sh`, e.g. `--sims_per_file 1 --parallel_jobs 4`
@@ -74,17 +81,54 @@ Notes:
 
 ```
 OpenAlphaDiffract/
-├── configs/ - Pipeline configuration files
-├── docker/ - Container definitions
-├── scripts/ - User-facing scripts
-├── src/ - Source code for pipeline components
-│   ├── downloader/
-│   ├── simulator/
-│   ├── trainer/
-│   └── ui/
+├── configs/              - Pipeline configuration files
+│   ├── instruments/      - GSAS-II instrument parameter files
+│   └── resources/        - Space group distance matrix for GEMD loss
+├── docker/               - Container definitions
+├── scripts/              - User-facing scripts
+├── src/
+│   ├── downloader/       - Materials Project CIF acquisition
+│   ├── simulator/        - GSAS-II powder diffraction simulation
+│   ├── trainer/          - Model definition, dataset, and training loop
+│   └── ui/               - FastAPI backend + React frontend
+│       └── frontend/
+└── compose.yaml
 ```
 
 
+## Testing
+
+Tests run in CI via GitHub Actions on every push/PR to `main`. To run locally:
+
+```bash
+# All Python tests (from repo root)
+pytest
+
+# Individual components
+pytest src/downloader/tests/ -v
+pytest src/simulator/tests/ -v
+pytest src/trainer/tests/ -v
+pytest src/ui/tests/ -v
+
+# Frontend tests
+cd src/ui/frontend && npx vitest run
+```
+
+
+## License
+
+This project is licensed under the BSD 3-Clause License. See [LICENSE](LICENSE) for details.
+
+
 ## Citation
 
-We hope this code was helpful to your work! If you use our code or extend our work, please consider citing our paper. (Will be added once released!)
+We hope this code was helpful! Please consider citing our paper:
+
+```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.},
+  journal={arXiv preprint arXiv:2603.23367},
+  year={2026}
+}
+```