Skip to content

Commit 070c2cf

Browse files
fallettafalletta
committed
dielectric constant SiO2 with allegro-pol
Made-with: Cursor
1 parent 8840c5e commit 070c2cf

5 files changed

Lines changed: 265 additions & 37 deletions

File tree

‎examples/allegro-pol/SiO2.xyz‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
72
2+
Lattice="9.97060582 0.0 0.0 -4.98530291 8.634797931240975 0.0 0.0 0.0 10.95705942" Properties=species:S:1:pos:R:3 pbc="T T T"
3+
Si -1.179706301 2.043311228 4.565438736
4+
Si 1.312964285 2.274118001 2.739272561
5+
Si 2.359415637 0.000002954 0.913100053
6+
O -1.083687128 3.188843421 0.237177664
7+
O 0.810784832 3.661485255 2.063349026
8+
O 0.272872565 1.784530781 3.889529133
9+
O 1.409015610 1.128571505 1.588994085
10+
O 2.765537669 2.532905391 3.415172492
11+
O 3.303456879 0.655920717 5.241349682
12+
Si -1.179705260 2.043315427 10.043959527
13+
Si 1.312961838 2.274119413 8.217787079
14+
Si 2.359412522 -0.000000046 6.391620844
15+
O -1.083684459 3.188838044 5.715710107
16+
O 0.810791514 3.661478811 7.541887148
17+
O 0.272863734 1.784529054 9.368065793
18+
O 1.409012690 1.128580016 7.067530745
19+
O 2.765546591 2.532902826 8.893710614
20+
O 3.303462870 0.655921094 10.719882125
21+
Si -3.672340860 6.360685423 4.565441174
22+
Si -1.179687683 6.591499487 2.739281798
23+
Si -0.133261917 4.317404554 0.913096022
24+
O -3.576287140 7.506226145 0.237176449
25+
O -1.681852974 7.978880156 2.063372286
26+
O -2.219756816 6.101880477 3.889531188
27+
O -1.083669333 5.445975836 1.588991616
28+
O 0.272887531 6.850288293 3.415173964
29+
O 0.810786438 4.973326281 5.241358021
30+
Si -3.672340447 6.360685933 10.043966769
31+
Si -1.179693570 6.591496140 8.217795120
32+
Si -0.133266276 4.317401471 6.391621412
33+
O -3.576290075 7.506227071 5.715702937
34+
O -1.681848414 7.978867852 7.541904085
35+
O -2.219765886 6.101879123 9.368073101
36+
O -1.083685151 5.445974061 7.067531367
37+
O 0.272899166 6.850277697 8.893703134
38+
O 0.810792715 4.973320112 10.719881752
39+
Si 3.805612490 2.043288824 4.565438168
40+
Si 6.298249911 2.274103458 2.739264520
41+
Si 7.344685767 0.000002813 0.913092811
42+
O 3.901625177 3.188827252 0.237178037
43+
O 5.796064860 3.661476698 2.063356506
44+
O 5.258194406 1.784488149 3.889528511
45+
O 6.394265257 1.128565967 1.588986777
46+
O 7.750826221 2.532910412 3.415155555
47+
O 8.288728263 0.655955446 5.241356852
48+
Si 3.805612981 2.043294141 10.043963558
49+
Si 6.298249867 2.274110229 8.217777842
50+
Si 7.344685532 0.000002200 6.391618406
51+
O 3.901633658 3.188824901 5.715701768
52+
O 5.796079854 3.661471920 7.541885676
53+
O 5.258188035 1.784502735 9.368068262
54+
O 6.394261894 1.128574498 7.067528690
55+
O 7.750839156 2.532912614 8.893687354
56+
O 8.288725993 0.655957525 10.719883340
57+
Si 1.312961077 6.360683803 4.565452347
58+
Si 3.805599879 6.591495121 2.739275398
59+
Si 4.852027080 4.317401431 0.913093439
60+
O 1.408987308 7.506217018 0.237171207
61+
O 3.303412894 7.978863555 2.063359531
62+
O 2.765540050 6.101882464 3.889545809
63+
O 3.901610116 5.445960007 1.588985577
64+
O 5.258180124 6.850273932 3.415163351
65+
O 5.796087987 4.973325423 5.241351020
66+
Si 1.312965778 6.360679496 10.043966141
67+
Si 3.805602305 6.591493721 8.217784242
68+
Si 4.852033160 4.317399514 6.391607233
69+
O 1.408982140 7.506222259 5.715708769
70+
O 3.303421121 7.978854769 7.541896289
71+
O 2.765534602 6.101873469 9.368074301
72+
O 3.901615182 5.445969223 7.067514069
73+
O 5.258191846 6.850271200 8.893700109
74+
O 5.796080865 4.973327277 10.719888582

‎examples/allegro-pol/allegro_pol_sio2.py‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
"""Relax alpha-quartz SiO2 under a uniform electric field with allegro-pol.
2+
3+
Initial geometry is read from ``SiO2-sc222.xyz``.
4+
5+
Two-stage relaxation (L-BFGS):
6+
1. Relax at E=0 to find the equilibrium geometry and bare polarization P0.
7+
2. Relax at E≠0 from that geometry and track the dielectric constant.
8+
"""
9+
10+
from __future__ import annotations
11+
12+
from pathlib import Path
13+
14+
import numpy as np
15+
import plotly.graph_objects as go
16+
import torch
17+
from allegro_pol.integrations.torchsim import NequIPPolTorchSimCalc
18+
from ase.io import read as ase_read
19+
from pymatgen.io.ase import AseAtomsAdaptor
20+
21+
import torch_sim as ts
22+
from torch_sim.models.interface import SerialSumModel
23+
from torch_sim.models.polarization import UniformPolarizationModel
24+
25+
26+
MODEL_PATH = "allegro-pol-SiO2.nequip.pt2"
27+
STRUCTURE_XYZ_PATH = Path(__file__).resolve().parent / "SiO2.xyz"
28+
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
29+
DTYPE = torch.float64
30+
31+
EV_PER_ANG2_TO_INV_EPSILON0 = 5.5263499562e-3
32+
33+
MAX_STEPS = 500
34+
FMAX_TOL = 1e-4
35+
E_FIELD_V_PER_ANG = 1e-3 # V/Å along z
36+
37+
PLOT_HTML_PATH = "eps_zz_vs_step.html"
38+
39+
atoms = ase_read(STRUCTURE_XYZ_PATH)
40+
structure = AseAtomsAdaptor.get_structure(atoms)
41+
42+
allegro_model = NequIPPolTorchSimCalc.from_compiled_model(
43+
MODEL_PATH, device=DEVICE, chemical_species_to_atom_type_map=True
44+
)
45+
46+
field_model = UniformPolarizationModel(
47+
device=allegro_model.device,
48+
dtype=allegro_model.dtype,
49+
compute_forces=True,
50+
compute_stress=True,
51+
)
52+
53+
combined_model = SerialSumModel(
54+
allegro_model,
55+
field_model,
56+
key_map={
57+
"polarization": "total_polarization",
58+
"born_charges": "born_effective_charges",
59+
},
60+
)
61+
62+
print("Relaxing structure at Efield = [0.0, 0.0, 0.0]")
63+
state_e0 = ts.initialize_state([structure], combined_model.device, combined_model.dtype)
64+
state_e0._system_extras["external_E_field"] = torch.zeros(
65+
1, 3, device=state_e0.device, dtype=state_e0.dtype
66+
)
67+
68+
optim_e0 = ts.lbfgs_init(state=state_e0, model=combined_model)
69+
for step in range(1, MAX_STEPS + 1):
70+
optim_e0 = ts.lbfgs_step(state=optim_e0, model=combined_model)
71+
fmax = optim_e0.forces.abs().max().cpu().item()
72+
if fmax < FMAX_TOL:
73+
break
74+
75+
P0 = optim_e0.total_polarization.clone()
76+
V = optim_e0.volume
77+
78+
alpha = optim_e0.polarizability[0].cpu().numpy()
79+
V0 = V[0].cpu().item()
80+
eps_electronic = np.eye(3) + alpha / (EV_PER_ANG2_TO_INV_EPSILON0 * V0)
81+
82+
print(f"Relaxing structure at Efield = [0.0, 0.0, {E_FIELD_V_PER_ANG}]")
83+
state_ef = ts.SimState(
84+
positions=optim_e0.positions.clone(),
85+
masses=optim_e0.masses.clone(),
86+
cell=optim_e0.cell.clone(),
87+
pbc=optim_e0.pbc,
88+
atomic_numbers=optim_e0.atomic_numbers.clone(),
89+
system_idx=optim_e0.system_idx.clone(),
90+
)
91+
state_ef._system_extras["external_E_field"] = torch.tensor(
92+
[[0.0, 0.0, E_FIELD_V_PER_ANG]], device=state_ef.device, dtype=state_ef.dtype
93+
)
94+
95+
optim_ef = ts.lbfgs_init(state=state_ef, model=combined_model, alpha=70.0, step_size=1.0)
96+
97+
P0_z = P0[0, 2].item()
98+
99+
pol_history = [optim_ef.total_polarization.cpu().numpy().copy()]
100+
V_history = [optim_ef.volume[0].cpu().item()]
101+
102+
for step in range(1, MAX_STEPS + 1):
103+
optim_ef = ts.lbfgs_step(state=optim_ef, model=combined_model)
104+
pol_history.append(optim_ef.total_polarization.cpu().numpy().copy())
105+
V_history.append(optim_ef.volume[0].cpu().item())
106+
fmax = optim_ef.forces.abs().max().cpu().item()
107+
if fmax < FMAX_TOL:
108+
break
109+
110+
pol_history = np.array(pol_history)
111+
V_history = np.array(V_history)
112+
113+
denom_t = E_FIELD_V_PER_ANG * EV_PER_ANG2_TO_INV_EPSILON0 * V_history
114+
115+
eps_static_z = 1.0 + (pol_history[:, 0, 2] - P0_z) / denom_t
116+
117+
steps = np.arange(len(eps_static_z))
118+
fig = go.Figure()
119+
120+
fig.add_trace(
121+
go.Scatter(
122+
x=steps,
123+
y=eps_static_z,
124+
mode="lines+markers",
125+
marker=dict(size=4),
126+
name="ε_0,zz",
127+
line=dict(color="#1f77b4"),
128+
),
129+
)
130+
fig.add_hline(
131+
y=eps_electronic[2, 2],
132+
line_dash="dot",
133+
line_color="gray",
134+
annotation_text=f"ε_∞,zz = {eps_electronic[2, 2]:.2f}",
135+
)
136+
137+
fig.update_xaxes(title_text="Relaxation step")
138+
fig.update_yaxes(title_text="ε_0,zz")
139+
fig.update_layout(title="ε_0,zz vs relaxation step")
140+
fig.write_html(PLOT_HTML_PATH)
141+
eps_inf = float(eps_electronic[2, 2])
142+
eps_0 = float(eps_static_z[-1])
143+
print(f"eps_inf = {eps_inf}")
144+
print(f"eps_0 = {eps_0}")
145+
fig.show()

‎pyproject.toml‎

Lines changed: 18 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[project]
22
name = "torch-sim-atomistic"
3-
version = "0.6.0"
3+
version = "0.5.2"
44
description = "A pytorch toolkit for calculating material properties using MLIPs"
55
authors = [
66
{ name = "Abhijeet Gangan", email = "abhijeetgangan@g.ucla.edu" },
@@ -39,7 +39,6 @@ dependencies = [
3939
[project.optional-dependencies]
4040
test = [
4141
"torch-sim-atomistic[io,symmetry,vesin]",
42-
"physical-validation>=1.0.5",
4342
"platformdirs>=4.0.0",
4443
"psutil>=7.0.0",
4544
"pymatgen>=2025.6.14",
@@ -58,6 +57,7 @@ orb = ["orb-models>=0.6.2"]
5857
sevenn = ["sevenn[torchsim]>=0.12.1"]
5958
graphpes = ["graph-pes>=0.1", "mace-torch>=0.3.12"]
6059
nequip = ["nequip>=0.17.1"]
60+
allegro_pol = ["allegro-pol>=0.1.0"]
6161
nequix = ["nequix[torch-sim]>=0.4.5"]
6262
fairchem = ["fairchem-core>=2.7", "scipy<1.17.0"]
6363
docs = [
@@ -88,6 +88,9 @@ build-backend = "uv_build"
8888
module-name = "torch_sim"
8989
module-root = ""
9090

91+
[tool.uv.sources]
92+
allegro-pol = { git = "https://github.com/mir-group/allegro-pol" }
93+
9194
[tool.ruff]
9295
target-version = "py312"
9396
line-length = 90
@@ -128,25 +131,20 @@ isort.lines-after-imports = 2
128131
pep8-naming.ignore-names = ["get_kT", "kT"]
129132

130133
[tool.ruff.lint.per-file-ignores]
131-
"**/tests/*" = ["ANN001", "ANN201", "ANN202", "D", "INP001", "S101", "T201"]
132-
"examples/**/*" = ["ANN001", "ANN201", "ANN202", "B018", "T201"]
134+
"**/tests/*" = ["ANN201", "D", "INP001", "S101"]
135+
"examples/**/*" = ["B018", "T201"]
133136
"examples/tutorials/**/*" = ["ALL"]
134-
"docs/tutorials/*.ipynb" = ["ANN001", "ANN201", "B905", "BLE001", "E501", "F841", "N816", "PLR1714", "RUF001", "T201"]
135137

136138
[tool.ruff.format]
137139
docstring-code-format = true
138140

139141
[tool.codespell]
140142
check-filenames = true
141143
ignore-words-list = ["convertor"] # codespell:ignore convertor
142-
skip = "docs/tutorials/integrator_tests_analysis.ipynb"
143144

144145
[tool.pytest.ini_options]
145-
addopts = ["-p no:warnings", "-m not physical_validation"]
146+
addopts = ["-p no:warnings"]
146147
testpaths = ["tests"]
147-
markers = [
148-
"physical_validation: long-running physical validation tests (run with: pytest -m physical_validation)",
149-
]
150148

151149
[tool.uv]
152150
# make these dependencies mutually exclusive since they use incompatible e3nn versions
@@ -156,14 +154,6 @@ conflicts = [
156154
{ extra = "fairchem" },
157155
{ extra = "graphpes" },
158156
],
159-
[
160-
{ extra = "fairchem" },
161-
{ extra = "graphpes" },
162-
],
163-
[
164-
{ extra = "fairchem" },
165-
{ extra = "mace" },
166-
],
167157
[
168158
{ extra = "fairchem" },
169159
{ extra = "mace" },
@@ -184,6 +174,10 @@ conflicts = [
184174
{ extra = "graphpes" },
185175
{ extra = "sevenn" },
186176
],
177+
[
178+
{ extra = "graphpes" },
179+
{ extra = "allegro_pol" },
180+
],
187181
[
188182
{ extra = "mace" },
189183
{ extra = "mattersim" },
@@ -200,6 +194,10 @@ conflicts = [
200194
{ extra = "mace" },
201195
{ extra = "sevenn" },
202196
],
197+
[
198+
{ extra = "mace" },
199+
{ extra = "allegro_pol" },
200+
],
203201
]
204202

205203
[dependency-groups]
@@ -226,36 +224,22 @@ include = [
226224
[tool.ty.overrides.rules]
227225
unresolved-import = "ignore"
228226

229-
[[tool.ty.overrides]]
230-
include = [
231-
"torch_sim/models/dispersion.py",
232-
"torch_sim/neighbors/vesin.py",
233-
]
234-
[tool.ty.overrides.rules]
235-
invalid-argument-type = "ignore"
236-
invalid-assignment = "ignore"
237-
238227
[[tool.ty.overrides]]
239228
include = ["tests/**/*.py"]
240229

241230
[tool.ty.overrides.rules]
242231
invalid-argument-type = "ignore"
243-
invalid-assignment = "ignore"
244232
no-matching-overload = "ignore"
245233
unresolved-attribute = "ignore"
246234
unresolved-import = "ignore"
247235

248236
[[tool.ty.overrides]]
249-
include = ["docs/**/*.py", "docs/**/*.ipynb", "examples/**/*.py"]
237+
include = ["docs/**/*.py", "examples/**/*.py"]
250238
[tool.ty.overrides.rules]
251-
invalid-argument-type = "ignore"
252-
not-iterable = "ignore"
253-
not-subscriptable = "ignore"
254-
unresolved-attribute = "ignore"
255239
unresolved-import = "ignore"
256240

257241

258242
[[tool.ty.overrides]]
259243
include = ["torch_sim/neighbors/alchemiops.py"]
260244
[tool.ty.overrides.rules]
261-
call-non-callable = "ignore"
245+
call-non-callable = "ignore"

0 commit comments

Comments
 (0)