File size: 73,667 Bytes
f0fdb88
 
 
 
b22d49d
f0fdb88
b22d49d
 
 
 
 
 
 
 
 
f0fdb88
 
 
 
 
 
 
 
 
 
 
 
 
c1d2139
 
1effaf5
 
 
 
 
 
8b028aa
1effaf5
 
 
8b028aa
f0fdb88
 
 
 
c7922c2
f0fdb88
 
 
 
 
 
 
 
 
 
 
 
 
 
c7922c2
f0fdb88
 
 
 
 
 
 
 
 
 
c7922c2
f0fdb88
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c7922c2
f0fdb88
c7922c2
 
 
 
 
 
 
 
 
f0fdb88
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1effaf5
 
 
 
 
 
 
f0fdb88
 
b22d49d
f0fdb88
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b22d49d
f0fdb88
 
 
 
 
 
 
 
 
 
c1d2139
 
 
f0fdb88
 
c1d2139
 
f0fdb88
 
 
8b028aa
f0fdb88
 
b22d49d
 
 
 
 
8b028aa
b22d49d
8b028aa
b22d49d
 
f0fdb88
 
 
 
 
 
 
c1d2139
 
 
 
 
 
 
 
f0fdb88
 
 
 
c1d2139
 
 
 
 
 
 
 
 
 
 
 
f0fdb88
1effaf5
c1d2139
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f0fdb88
1effaf5
f0fdb88
 
 
 
 
b22d49d
f0fdb88
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c7922c2
f0fdb88
cf512f3
f0fdb88
 
 
 
 
 
 
 
 
 
 
b22d49d
f0fdb88
 
 
b22d49d
 
8b028aa
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b22d49d
 
 
8b028aa
b22d49d
8b028aa
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b22d49d
 
 
f0fdb88
 
 
 
 
 
 
 
 
c1d2139
 
 
f0fdb88
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8b028aa
f0fdb88
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No module named 'deepmd'\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "from ase import units, Atoms\n",
    "from ase.build import molecule\n",
    "from ase.io import read, write\n",
    "from dask.distributed import Client\n",
    "from dask_jobqueue import SLURMCluster\n",
    "from prefect import flow\n",
    "from prefect_dask import DaskTaskRunner\n",
    "from pymatgen.core import Molecule\n",
    "from pymatgen.io.packmol import PackmolBoxGen\n",
    "\n",
    "from mlip_arena.models import REGISTRY, MLIPEnum\n",
    "\n",
    "from mlip_arena.tasks.md import run as MD"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Intial configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "h2 = molecule(\"H2\")\n",
    "o2 = molecule(\"O2\")\n",
    "h2o = molecule(\"H2O\")\n",
    "\n",
    "write(\"h2.xyz\", h2)\n",
    "write(\"o2.xyz\", o2)\n",
    "write(\"h2o.xyz\", h2o)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "h2 = Molecule.from_file(\"h2.xyz\")\n",
    "o2 = Molecule.from_file(\"o2.xyz\")\n",
    "h2o = Molecule.from_file(\"h2o.xyz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "molecules = []\n",
    "\n",
    "for m, number in zip([h2, o2], [128, 64]):\n",
    "    molecules.append(\n",
    "        {\n",
    "            \"name\": m.composition.to_pretty_string(),\n",
    "            \"number\": number,\n",
    "            \"coords\": m,\n",
    "        }\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Atoms(symbols='H256O128', pbc=True, cell=[30.0, 30.0, 30.0])\n"
     ]
    }
   ],
   "source": [
    "tolerance = 2.0\n",
    "input_gen = PackmolBoxGen(\n",
    "    tolerance=tolerance,\n",
    "    seed=1,\n",
    ")\n",
    "margin = 0.5 * tolerance\n",
    "\n",
    "a = 30\n",
    "\n",
    "packmol_set = input_gen.get_input_set(\n",
    "    molecules=molecules,\n",
    "    box=[margin, margin, margin, a - margin, a - margin, a - margin],\n",
    ")\n",
    "packmol_set.write_input(\".\")\n",
    "packmol_set.run(\".\")\n",
    "\n",
    "atoms = read(\"packmol_out.xyz\")\n",
    "atoms.cell = [a, a, a]\n",
    "atoms.pbc = True\n",
    "\n",
    "print(atoms)\n",
    "\n",
    "write(f'{atoms.get_chemical_formula()}.extxyz', atoms)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Atoms(symbols='H256O128', pbc=True, cell=[30.0, 30.0, 30.0])\n"
     ]
    }
   ],
   "source": [
    "atoms = read(\"H256O128.extxyz\")\n",
    "print(atoms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#!/bin/bash\n",
      "\n",
      "#SBATCH -A matgen\n",
      "#SBATCH --mem=0\n",
      "#SBATCH -t 02:00:00\n",
      "#SBATCH -J combustion-water\n",
      "#SBATCH -q regular\n",
      "#SBATCH -N 1\n",
      "#SBATCH -C gpu\n",
      "#SBATCH -G 4\n",
      "#SBATCH --exclusive\n",
      "source ~/.bashrc\n",
      "module load python\n",
      "source activate /pscratch/sd/c/cyrusyc/.conda/mlip-arena\n",
      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/bin/python -m distributed.cli.dask_worker tcp://128.55.64.15:38781 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 86400\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/node.py:187: UserWarning: Port 8787 is already in use.\n",
      "Perhaps you already have a cluster running?\n",
      "Hosting the HTTP server on port 44831 instead\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "nodes_per_alloc = 1\n",
    "gpus_per_alloc = 4\n",
    "ntasks = 1\n",
    "\n",
    "cluster_kwargs = dict(\n",
    "    cores=1,\n",
    "    memory=\"64 GB\",\n",
    "    shebang=\"#!/bin/bash\",\n",
    "    account=\"matgen\",\n",
    "    walltime=\"02:00:00\",\n",
    "    job_mem=\"0\",\n",
    "    job_script_prologue=[\n",
    "        \"source ~/.bashrc\",\n",
    "        \"module load python\",\n",
    "        \"source activate /pscratch/sd/c/cyrusyc/.conda/mlip-arena\",\n",
    "    ],\n",
    "    job_directives_skip=[\"-n\", \"--cpus-per-task\", \"-J\"],\n",
    "    job_extra_directives=[\n",
    "        \"-J combustion-water\",\n",
    "        \"-q regular\",\n",
    "        f\"-N {nodes_per_alloc}\",\n",
    "        \"-C gpu\",\n",
    "        f\"-G {gpus_per_alloc}\",\n",
    "        f\"--exclusive\",\n",
    "    ],\n",
    "    death_timeout=86400\n",
    ")\n",
    "\n",
    "cluster = SLURMCluster(**cluster_kwargs)\n",
    "\n",
    "# cluster_kwargs = {\n",
    "#     \"cores\": 1,\n",
    "#     \"memory\": \"64 GB\",\n",
    "#     \"shebang\": \"#!/bin/bash\",\n",
    "#     \"account\": \"matgen\",\n",
    "#     \"walltime\": \"00:30:00\",\n",
    "#     \"job_mem\": \"0\",\n",
    "#     \"job_script_prologue\": [\n",
    "#         \"source ~/.bashrc\",\n",
    "#         \"module load python\",\n",
    "#         \"source activate /pscratch/sd/c/cyrusyc/.conda/mlip-arena\",\n",
    "#     ],\n",
    "#     \"job_directives_skip\": [\"-n\", \"--cpus-per-task\", \"-J\"],\n",
    "#     \"job_extra_directives\": [f\"-N {nodes_per_alloc}\", \"-q debug\", \"-C gpu\", \"-J combustion-water\"],\n",
    "# }\n",
    "# cluster = SLURMCluster(**cluster_kwargs)\n",
    "\n",
    "print(cluster.job_script())\n",
    "cluster.adapt(minimum_jobs=2, maximum_jobs=2)\n",
    "client = Client(cluster)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "@flow(task_runner=DaskTaskRunner(address=client.scheduler.address), log_prints=True)\n",
    "def combustion(atoms: Atoms):\n",
    "    futures = []\n",
    "\n",
    "    for model in MLIPEnum:\n",
    "        future = MD.submit(\n",
    "            atoms=atoms,\n",
    "            calculator_name=model,\n",
    "            calculator_kwargs=None,\n",
    "            ensemble=\"nvt\",\n",
    "            dynamics=\"nose-hoover\",\n",
    "            time_step=None,\n",
    "            ase_md_kwargs=dict(ttime=25 * units.fs, pfactor=None),\n",
    "            total_time=1000_000,\n",
    "            temperature=[300, 3000, 3000, 300],\n",
    "            pressure=None,\n",
    "            md_velocity_seed=0,\n",
    "            traj_file=Path(REGISTRY[model.name][\"family\"])\n",
    "            / f\"{model.name}_{atoms.get_chemical_formula()}.traj\",\n",
    "            traj_interval=1000,\n",
    "            restart=True,\n",
    "        )\n",
    "\n",
    "        futures.append(future)\n",
    "    \n",
    "    return [future.result() for future in futures]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">12:01:53.150 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - Created flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'khaki-hippo'</span> for flow<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> 'combustion'</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "12:01:53.150 | \u001b[36mINFO\u001b[0m    | prefect.engine - Created flow run\u001b[35m 'khaki-hippo'\u001b[0m for flow\u001b[1;35m 'combustion'\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">12:01:53.156 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - View at <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/runs/flow-run/ce836160-4f0d-49a4-90d8-227225cb8f4c</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "12:01:53.156 | \u001b[36mINFO\u001b[0m    | prefect.engine - View at \u001b[94mhttps://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/runs/flow-run/ce836160-4f0d-49a4-90d8-227225cb8f4c\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">12:01:53.523 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(99282258, 'tcp://128.55.64.15:38781', workers=0, threads=0, memory=0 B)\n",
       "</pre>\n"
      ],
      "text/plain": [
       "12:01:53.523 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(99282258, 'tcp://128.55.64.15:38781', workers=0, threads=0, memory=0 B)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">13:06:57.332 | <span style=\"color: #d70000; text-decoration-color: #d70000\">ERROR</span>   | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'khaki-hippo'</span> - Encountered exception during execution: IndexError('too many indices for tensor of dimension 1')\n",
       "Traceback (most recent call last):\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 652, in run_context\n",
       "    yield self\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 696, in run_flow_sync\n",
       "    engine.call_flow_fn()\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 675, in call_flow_fn\n",
       "    result = call_with_parameters(self.flow.fn, self.parameters)\n",
       "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py\", line 206, in call_with_parameters\n",
       "    return fn(*args, **kwargs)\n",
       "           ^^^^^^^^^^^^^^^^^^^\n",
       "  File \"/tmp/ipykernel_1849247/2043615938.py\", line 26, in combustion\n",
       "    return [future.result() for future in futures]\n",
       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
       "  File \"/tmp/ipykernel_1849247/2043615938.py\", line 26, in &lt;listcomp&gt;\n",
       "    return [future.result() for future in futures]\n",
       "            ^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/task_runners.py\", line 143, in result\n",
       "    _result = self._final_state.result(\n",
       "              ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/client/schemas/objects.py\", line 314, in result\n",
       "    return get_state_result(\n",
       "           ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/states.py\", line 75, in get_state_result\n",
       "    return _get_state_result(\n",
       "           ^^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/asyncutils.py\", line 399, in coroutine_wrapper\n",
       "    return run_coro_as_sync(ctx_call())\n",
       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/asyncutils.py\", line 243, in run_coro_as_sync\n",
       "    return call.result()\n",
       "           ^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 312, in result\n",
       "    return self.future.result(timeout=timeout)\n",
       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 182, in result\n",
       "    return self.__get_result()\n",
       "           ^^^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/concurrent/futures/_base.py\", line 401, in __get_result\n",
       "    raise self._exception\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 383, in _run_async\n",
       "    result = await coro\n",
       "             ^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/asyncutils.py\", line 225, in coroutine_wrapper\n",
       "    return await task\n",
       "           ^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/asyncutils.py\", line 389, in ctx_call\n",
       "    result = await async_fn(*args, **kwargs)\n",
       "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/states.py\", line 138, in _get_state_result\n",
       "    raise await get_state_exception(state)\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/task_engine.py\", line 763, in run_context\n",
       "    yield self\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/task_engine.py\", line 1323, in run_task_sync\n",
       "    engine.call_task_fn(txn)\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/task_engine.py\", line 786, in call_task_fn\n",
       "    result = call_with_parameters(self.task.fn, parameters)\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py\", line 206, in call_with_parameters\n",
       "    return fn(*args, **kwargs)\n",
       "      ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/mlip-arena/mlip_arena/tasks/md.py\", line 363, in run\n",
       "    md_runner.run(steps=n_steps)\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ase/md/npt.py\", line 277, in run\n",
       "    self.initialize()\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ase/md/npt.py\", line 399, in initialize\n",
       "    self._calculate_q_past_and_future()\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ase/md/npt.py\", line 608, in _calculate_q_past_and_future\n",
       "    self._calculate_q_future(self.atoms.get_forces(md=True))\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ase/atoms.py\", line 812, in get_forces\n",
       "    forces = self._calc.get_forces(self)\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ase/calculators/abc.py\", line 30, in get_forces\n",
       "    return self.get_property('forces', atoms)\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ase/calculators/calculator.py\", line 538, in get_property\n",
       "    self.calculate(atoms, [name], system_changes)\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/mlip-arena/mlip_arena/models/externals/chgnet.py\", line 36, in calculate\n",
       "    super().calculate(atoms, properties, system_changes)\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/chgnet/model/dynamics.py\", line 143, in calculate\n",
       "    model_prediction = self.model.predict_graph(\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/chgnet/model/model.py\", line 627, in predict_graph\n",
       "    prediction = self.forward(\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/chgnet/model/model.py\", line 359, in forward\n",
       "    batched_graph = BatchedGraph.from_graphs(\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/chgnet/model/model.py\", line 822, in from_graphs\n",
       "    center=atom_cart_coords[graph.atom_graph[:, 0]],\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "IndexError: too many indices for tensor of dimension 1\n",
       "</pre>\n"
      ],
      "text/plain": [
       "13:06:57.332 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'khaki-hippo'\u001b[0m - Encountered exception during execution: IndexError('too many indices for tensor of dimension 1')\n",
       "Traceback (most recent call last):\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 652, in run_context\n",
       "    yield self\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 696, in run_flow_sync\n",
       "    engine.call_flow_fn()\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 675, in call_flow_fn\n",
       "    result = call_with_parameters(self.flow.fn, self.parameters)\n",
       "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py\", line 206, in call_with_parameters\n",
       "    return fn(*args, **kwargs)\n",
       "           ^^^^^^^^^^^^^^^^^^^\n",
       "  File \"/tmp/ipykernel_1849247/2043615938.py\", line 26, in combustion\n",
       "    return [future.result() for future in futures]\n",
       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
       "  File \"/tmp/ipykernel_1849247/2043615938.py\", line 26, in <listcomp>\n",
       "    return [future.result() for future in futures]\n",
       "            ^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/task_runners.py\", line 143, in result\n",
       "    _result = self._final_state.result(\n",
       "              ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/client/schemas/objects.py\", line 314, in result\n",
       "    return get_state_result(\n",
       "           ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/states.py\", line 75, in get_state_result\n",
       "    return _get_state_result(\n",
       "           ^^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/asyncutils.py\", line 399, in coroutine_wrapper\n",
       "    return run_coro_as_sync(ctx_call())\n",
       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/asyncutils.py\", line 243, in run_coro_as_sync\n",
       "    return call.result()\n",
       "           ^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 312, in result\n",
       "    return self.future.result(timeout=timeout)\n",
       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 182, in result\n",
       "    return self.__get_result()\n",
       "           ^^^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/concurrent/futures/_base.py\", line 401, in __get_result\n",
       "    raise self._exception\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 383, in _run_async\n",
       "    result = await coro\n",
       "             ^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/asyncutils.py\", line 225, in coroutine_wrapper\n",
       "    return await task\n",
       "           ^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/asyncutils.py\", line 389, in ctx_call\n",
       "    result = await async_fn(*args, **kwargs)\n",
       "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/states.py\", line 138, in _get_state_result\n",
       "    raise await get_state_exception(state)\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/task_engine.py\", line 763, in run_context\n",
       "    yield self\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/task_engine.py\", line 1323, in run_task_sync\n",
       "    engine.call_task_fn(txn)\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/task_engine.py\", line 786, in call_task_fn\n",
       "    result = call_with_parameters(self.task.fn, parameters)\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py\", line 206, in call_with_parameters\n",
       "    return fn(*args, **kwargs)\n",
       "      ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/mlip-arena/mlip_arena/tasks/md.py\", line 363, in run\n",
       "    md_runner.run(steps=n_steps)\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ase/md/npt.py\", line 277, in run\n",
       "    self.initialize()\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ase/md/npt.py\", line 399, in initialize\n",
       "    self._calculate_q_past_and_future()\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ase/md/npt.py\", line 608, in _calculate_q_past_and_future\n",
       "    self._calculate_q_future(self.atoms.get_forces(md=True))\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ase/atoms.py\", line 812, in get_forces\n",
       "    forces = self._calc.get_forces(self)\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ase/calculators/abc.py\", line 30, in get_forces\n",
       "    return self.get_property('forces', atoms)\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ase/calculators/calculator.py\", line 538, in get_property\n",
       "    self.calculate(atoms, [name], system_changes)\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/mlip-arena/mlip_arena/models/externals/chgnet.py\", line 36, in calculate\n",
       "    super().calculate(atoms, properties, system_changes)\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/chgnet/model/dynamics.py\", line 143, in calculate\n",
       "    model_prediction = self.model.predict_graph(\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/chgnet/model/model.py\", line 627, in predict_graph\n",
       "    prediction = self.forward(\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/chgnet/model/model.py\", line 359, in forward\n",
       "    batched_graph = BatchedGraph.from_graphs(\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/chgnet/model/model.py\", line 822, in from_graphs\n",
       "    center=atom_cart_coords[graph.atom_graph[:, 0]],\n",
       "    ^^^^^^^^^^^^^^^^^\n",
       "IndexError: too many indices for tensor of dimension 1\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">13:06:58.009 | <span style=\"color: #d70000; text-decoration-color: #d70000\">ERROR</span>   | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'khaki-hippo'</span> - Finished in state <span style=\"color: #d70000; text-decoration-color: #d70000\">Failed</span>('Flow run encountered an exception: IndexError: too many indices for tensor of dimension 1')\n",
       "</pre>\n"
      ],
      "text/plain": [
       "13:06:58.009 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'khaki-hippo'\u001b[0m - Finished in state \u001b[38;5;160mFailed\u001b[0m('Flow run encountered an exception: IndexError: too many indices for tensor of dimension 1')\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "IndexError",
     "evalue": "too many indices for tensor of dimension 1",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[5], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m results \u001b[38;5;241m=\u001b[39m \u001b[43mcombustion\u001b[49m\u001b[43m(\u001b[49m\u001b[43matoms\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flows.py:1345\u001b[0m, in \u001b[0;36mFlow.__call__\u001b[0;34m(self, return_state, wait_for, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1341\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m track_viz_task(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39misasync, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mname, parameters)\n\u001b[1;32m   1343\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mprefect\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mflow_engine\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m run_flow\n\u001b[0;32m-> 1345\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrun_flow\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1346\u001b[0m \u001b[43m    \u001b[49m\u001b[43mflow\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1347\u001b[0m \u001b[43m    \u001b[49m\u001b[43mparameters\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mparameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1348\u001b[0m \u001b[43m    \u001b[49m\u001b[43mwait_for\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwait_for\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1349\u001b[0m \u001b[43m    \u001b[49m\u001b[43mreturn_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_type\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1350\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:818\u001b[0m, in \u001b[0;36mrun_flow\u001b[0;34m(flow, flow_run, parameters, wait_for, return_type)\u001b[0m\n\u001b[1;32m    816\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m run_flow_async(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    817\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 818\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrun_flow_sync\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:698\u001b[0m, in \u001b[0;36mrun_flow_sync\u001b[0;34m(flow, flow_run, parameters, wait_for, return_type)\u001b[0m\n\u001b[1;32m    695\u001b[0m         \u001b[38;5;28;01mwith\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mrun_context():\n\u001b[1;32m    696\u001b[0m             engine\u001b[38;5;241m.\u001b[39mcall_flow_fn()\n\u001b[0;32m--> 698\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mstate \u001b[38;5;28;01mif\u001b[39;00m return_type \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstate\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m \u001b[43mengine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:255\u001b[0m, in \u001b[0;36mFlowRunEngine.result\u001b[0;34m(self, raise_on_failure)\u001b[0m\n\u001b[1;32m    253\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_raised \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m NotSet:\n\u001b[1;32m    254\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m raise_on_failure:\n\u001b[0;32m--> 255\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_raised\n\u001b[1;32m    256\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_raised\n\u001b[1;32m    258\u001b[0m \u001b[38;5;66;03m# This is a fall through case which leans on the existing state result mechanics to get the\u001b[39;00m\n\u001b[1;32m    259\u001b[0m \u001b[38;5;66;03m# return value. This is necessary because we currently will return a State object if the\u001b[39;00m\n\u001b[1;32m    260\u001b[0m \u001b[38;5;66;03m# the State was Prefect-created.\u001b[39;00m\n\u001b[1;32m    261\u001b[0m \u001b[38;5;66;03m# TODO: Remove the need to get the result from a State except in cases where the return value\u001b[39;00m\n\u001b[1;32m    262\u001b[0m \u001b[38;5;66;03m# is a State object.\u001b[39;00m\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:652\u001b[0m, in \u001b[0;36mFlowRunEngine.run_context\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    645\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m timeout_context(\n\u001b[1;32m    646\u001b[0m         seconds\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mflow\u001b[38;5;241m.\u001b[39mtimeout_seconds,\n\u001b[1;32m    647\u001b[0m         timeout_exc_type\u001b[38;5;241m=\u001b[39mFlowRunTimeoutError,\n\u001b[1;32m    648\u001b[0m     ):\n\u001b[1;32m    649\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlogger\u001b[38;5;241m.\u001b[39mdebug(\n\u001b[1;32m    650\u001b[0m             \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mExecuting flow \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mflow\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m!r}\u001b[39;00m\u001b[38;5;124m for flow run \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mflow_run\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m!r}\u001b[39;00m\u001b[38;5;124m...\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    651\u001b[0m         )\n\u001b[0;32m--> 652\u001b[0m         \u001b[38;5;28;01myield\u001b[39;00m \u001b[38;5;28mself\u001b[39m\n\u001b[1;32m    653\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTimeoutError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[1;32m    654\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandle_timeout(exc)\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:696\u001b[0m, in \u001b[0;36mrun_flow_sync\u001b[0;34m(flow, flow_run, parameters, wait_for, return_type)\u001b[0m\n\u001b[1;32m    694\u001b[0m     \u001b[38;5;28;01mwhile\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mis_running():\n\u001b[1;32m    695\u001b[0m         \u001b[38;5;28;01mwith\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mrun_context():\n\u001b[0;32m--> 696\u001b[0m             \u001b[43mengine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcall_flow_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    698\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mstate \u001b[38;5;28;01mif\u001b[39;00m return_type \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstate\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mresult()\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:675\u001b[0m, in \u001b[0;36mFlowRunEngine.call_flow_fn\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    673\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m _call_flow_fn()\n\u001b[1;32m    674\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 675\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[43mcall_with_parameters\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mflow\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfn\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparameters\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    676\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandle_success(result)\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py:206\u001b[0m, in \u001b[0;36mcall_with_parameters\u001b[0;34m(fn, parameters)\u001b[0m\n\u001b[1;32m    198\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    199\u001b[0m \u001b[38;5;124;03mCall a function with parameters extracted with `get_call_parameters`\u001b[39;00m\n\u001b[1;32m    200\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    203\u001b[0m \u001b[38;5;124;03mthe args/kwargs using `parameters_to_positional_and_keyword` directly\u001b[39;00m\n\u001b[1;32m    204\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    205\u001b[0m args, kwargs \u001b[38;5;241m=\u001b[39m parameters_to_args_kwargs(fn, parameters)\n\u001b[0;32m--> 206\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[4], line 26\u001b[0m, in \u001b[0;36mcombustion\u001b[0;34m(atoms)\u001b[0m\n\u001b[1;32m      6\u001b[0m     future \u001b[38;5;241m=\u001b[39m MD\u001b[38;5;241m.\u001b[39msubmit(\n\u001b[1;32m      7\u001b[0m         atoms\u001b[38;5;241m=\u001b[39matoms,\n\u001b[1;32m      8\u001b[0m         calculator_name\u001b[38;5;241m=\u001b[39mmodel,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     21\u001b[0m         restart\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     22\u001b[0m     )\n\u001b[1;32m     24\u001b[0m     futures\u001b[38;5;241m.\u001b[39mappend(future)\n\u001b[0;32m---> 26\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m[\u001b[49m\u001b[43mfuture\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mfuture\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mfutures\u001b[49m\u001b[43m]\u001b[49m\n",
      "Cell \u001b[0;32mIn[4], line 26\u001b[0m, in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m      6\u001b[0m     future \u001b[38;5;241m=\u001b[39m MD\u001b[38;5;241m.\u001b[39msubmit(\n\u001b[1;32m      7\u001b[0m         atoms\u001b[38;5;241m=\u001b[39matoms,\n\u001b[1;32m      8\u001b[0m         calculator_name\u001b[38;5;241m=\u001b[39mmodel,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     21\u001b[0m         restart\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     22\u001b[0m     )\n\u001b[1;32m     24\u001b[0m     futures\u001b[38;5;241m.\u001b[39mappend(future)\n\u001b[0;32m---> 26\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m [\u001b[43mfuture\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m future \u001b[38;5;129;01min\u001b[39;00m futures]\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/task_runners.py:143\u001b[0m, in \u001b[0;36mPrefectDaskFuture.result\u001b[0;34m(self, timeout, raise_on_failure)\u001b[0m\n\u001b[1;32m    140\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    141\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m future_result\n\u001b[0;32m--> 143\u001b[0m _result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_final_state\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    144\u001b[0m \u001b[43m    \u001b[49m\u001b[43mraise_on_failure\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mraise_on_failure\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfetch\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\n\u001b[1;32m    145\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    146\u001b[0m \u001b[38;5;66;03m# state.result is a `sync_compatible` function that may or may not return an awaitable\u001b[39;00m\n\u001b[1;32m    147\u001b[0m \u001b[38;5;66;03m# depending on whether the parent frame is sync or not\u001b[39;00m\n\u001b[1;32m    148\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m asyncio\u001b[38;5;241m.\u001b[39miscoroutine(_result):\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/client/schemas/objects.py:314\u001b[0m, in \u001b[0;36mState.result\u001b[0;34m(self, raise_on_failure, fetch, retry_result_failure)\u001b[0m\n\u001b[1;32m    229\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    230\u001b[0m \u001b[38;5;124;03mRetrieve the result attached to this state.\u001b[39;00m\n\u001b[1;32m    231\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    310\u001b[0m \u001b[38;5;124;03m    >>> await flow_run.state.result(raise_on_failure=True, fetch=True) # Raises `ValueError(\"oh no!\")`\u001b[39;00m\n\u001b[1;32m    311\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    312\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mprefect\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mstates\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m get_state_result\n\u001b[0;32m--> 314\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mget_state_result\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    315\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m    316\u001b[0m \u001b[43m    \u001b[49m\u001b[43mraise_on_failure\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mraise_on_failure\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    317\u001b[0m \u001b[43m    \u001b[49m\u001b[43mfetch\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfetch\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    318\u001b[0m \u001b[43m    \u001b[49m\u001b[43mretry_result_failure\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mretry_result_failure\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    319\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/states.py:75\u001b[0m, in \u001b[0;36mget_state_result\u001b[0;34m(state, raise_on_failure, fetch, retry_result_failure)\u001b[0m\n\u001b[1;32m     73\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m state\u001b[38;5;241m.\u001b[39mdata\n\u001b[1;32m     74\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m---> 75\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_get_state_result\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     76\u001b[0m \u001b[43m        \u001b[49m\u001b[43mstate\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     77\u001b[0m \u001b[43m        \u001b[49m\u001b[43mraise_on_failure\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mraise_on_failure\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     78\u001b[0m \u001b[43m        \u001b[49m\u001b[43mretry_result_failure\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mretry_result_failure\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     79\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/asyncutils.py:399\u001b[0m, in \u001b[0;36msync_compatible.<locals>.coroutine_wrapper\u001b[0;34m(_sync, *args, **kwargs)\u001b[0m\n\u001b[1;32m    397\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m ctx_call()\n\u001b[1;32m    398\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 399\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrun_coro_as_sync\u001b[49m\u001b[43m(\u001b[49m\u001b[43mctx_call\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/asyncutils.py:243\u001b[0m, in \u001b[0;36mrun_coro_as_sync\u001b[0;34m(coroutine, force_new_thread, wait_for_result)\u001b[0m\n\u001b[1;32m    241\u001b[0m runner\u001b[38;5;241m.\u001b[39msubmit(call)\n\u001b[1;32m    242\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 243\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcall\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    244\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m:\n\u001b[1;32m    245\u001b[0m     call\u001b[38;5;241m.\u001b[39mcancel()\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:312\u001b[0m, in \u001b[0;36mCall.result\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    306\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mresult\u001b[39m(\u001b[38;5;28mself\u001b[39m, timeout: Optional[\u001b[38;5;28mfloat\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m T:\n\u001b[1;32m    307\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    308\u001b[0m \u001b[38;5;124;03m    Wait for the result of the call.\u001b[39;00m\n\u001b[1;32m    309\u001b[0m \n\u001b[1;32m    310\u001b[0m \u001b[38;5;124;03m    Not safe for use from asynchronous contexts.\u001b[39;00m\n\u001b[1;32m    311\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 312\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfuture\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:182\u001b[0m, in \u001b[0;36mFuture.result\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    180\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m CancelledError()\n\u001b[1;32m    181\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_state \u001b[38;5;241m==\u001b[39m FINISHED:\n\u001b[0;32m--> 182\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m__get_result\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    183\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    184\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTimeoutError\u001b[39;00m()\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/concurrent/futures/_base.py:401\u001b[0m, in \u001b[0;36mFuture.__get_result\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    399\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exception:\n\u001b[1;32m    400\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 401\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exception\n\u001b[1;32m    402\u001b[0m     \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    403\u001b[0m         \u001b[38;5;66;03m# Break a reference cycle with the exception in self._exception\u001b[39;00m\n\u001b[1;32m    404\u001b[0m         \u001b[38;5;28mself\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:383\u001b[0m, in \u001b[0;36mCall._run_async\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m    381\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfuture\u001b[38;5;241m.\u001b[39menforce_async_deadline() \u001b[38;5;28;01mas\u001b[39;00m cancel_scope:\n\u001b[1;32m    382\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 383\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m coro\n\u001b[1;32m    384\u001b[0m     \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    385\u001b[0m         \u001b[38;5;66;03m# Forget this call's arguments in order to free up any memory\u001b[39;00m\n\u001b[1;32m    386\u001b[0m         \u001b[38;5;66;03m# that may be referenced by them; after a call has happened,\u001b[39;00m\n\u001b[1;32m    387\u001b[0m         \u001b[38;5;66;03m# there's no need to keep a reference to them\u001b[39;00m\n\u001b[1;32m    388\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39margs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/asyncutils.py:225\u001b[0m, in \u001b[0;36mrun_coro_as_sync.<locals>.coroutine_wrapper\u001b[0;34m()\u001b[0m\n\u001b[1;32m    223\u001b[0m     task \u001b[38;5;241m=\u001b[39m create_task(coroutine)\n\u001b[1;32m    224\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m wait_for_result:\n\u001b[0;32m--> 225\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m task\n\u001b[1;32m    226\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    227\u001b[0m     RUNNING_IN_RUN_SYNC_LOOP_FLAG\u001b[38;5;241m.\u001b[39mreset(token1)\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/asyncutils.py:389\u001b[0m, in \u001b[0;36msync_compatible.<locals>.coroutine_wrapper.<locals>.ctx_call\u001b[0;34m()\u001b[0m\n\u001b[1;32m    387\u001b[0m token \u001b[38;5;241m=\u001b[39m RUNNING_ASYNC_FLAG\u001b[38;5;241m.\u001b[39mset(\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m    388\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 389\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m async_fn(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    390\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    391\u001b[0m     RUNNING_ASYNC_FLAG\u001b[38;5;241m.\u001b[39mreset(token)\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/states.py:138\u001b[0m, in \u001b[0;36m_get_state_result\u001b[0;34m(state, raise_on_failure, retry_result_failure)\u001b[0m\n\u001b[1;32m    131\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m UnfinishedRun(\n\u001b[1;32m    132\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mRun is in \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mstate\u001b[38;5;241m.\u001b[39mtype\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m state, its result is not available.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    133\u001b[0m     )\n\u001b[1;32m    135\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m raise_on_failure \u001b[38;5;129;01mand\u001b[39;00m (\n\u001b[1;32m    136\u001b[0m     state\u001b[38;5;241m.\u001b[39mis_crashed() \u001b[38;5;129;01mor\u001b[39;00m state\u001b[38;5;241m.\u001b[39mis_failed() \u001b[38;5;129;01mor\u001b[39;00m state\u001b[38;5;241m.\u001b[39mis_cancelled()\n\u001b[1;32m    137\u001b[0m ):\n\u001b[0;32m--> 138\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m get_state_exception(state)\n\u001b[1;32m    140\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(state\u001b[38;5;241m.\u001b[39mdata, (BaseResult, ResultRecordMetadata)):\n\u001b[1;32m    141\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m _get_state_result_data_with_retries(\n\u001b[1;32m    142\u001b[0m         state, retry_result_failure\u001b[38;5;241m=\u001b[39mretry_result_failure\n\u001b[1;32m    143\u001b[0m     )\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/task_engine.py:763\u001b[0m, in \u001b[0;36mrun_context\u001b[0;34m()\u001b[0m\n\u001b[1;32m    760\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mis_cancelled():\n\u001b[1;32m    761\u001b[0m             \u001b[38;5;28;01mraise\u001b[39;00m CancelledError(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTask run cancelled by the task runner\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m--> 763\u001b[0m         \u001b[38;5;28;01myield\u001b[39;00m \u001b[38;5;28mself\u001b[39m\n\u001b[1;32m    764\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTimeoutError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[1;32m    765\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandle_timeout(exc)\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/task_engine.py:1323\u001b[0m, in \u001b[0;36mrun_task_sync\u001b[0;34m()\u001b[0m\n\u001b[1;32m   1321\u001b[0m         run_coro_as_sync(engine\u001b[38;5;241m.\u001b[39mwait_until_ready())\n\u001b[1;32m   1322\u001b[0m         \u001b[38;5;28;01mwith\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mrun_context(), engine\u001b[38;5;241m.\u001b[39mtransaction_context() \u001b[38;5;28;01mas\u001b[39;00m txn:\n\u001b[0;32m-> 1323\u001b[0m             engine\u001b[38;5;241m.\u001b[39mcall_task_fn(txn)\n\u001b[1;32m   1325\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mstate \u001b[38;5;28;01mif\u001b[39;00m return_type \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstate\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mresult()\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/task_engine.py:786\u001b[0m, in \u001b[0;36mcall_task_fn\u001b[0;34m()\u001b[0m\n\u001b[1;32m    784\u001b[0m             result \u001b[38;5;241m=\u001b[39m call_with_parameters(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtask\u001b[38;5;241m.\u001b[39mfn, parameters)\n\u001b[1;32m    785\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 786\u001b[0m         result \u001b[38;5;241m=\u001b[39m call_with_parameters(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtask\u001b[38;5;241m.\u001b[39mfn, parameters)\n\u001b[1;32m    787\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandle_success(result, transaction\u001b[38;5;241m=\u001b[39mtransaction)\n\u001b[1;32m    788\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py:206\u001b[0m, in \u001b[0;36mcall_with_parameters\u001b[0;34m()\u001b[0m\n\u001b[1;32m    198\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    199\u001b[0m \u001b[38;5;124;03mCall a function with parameters extracted with `get_call_parameters`\u001b[39;00m\n\u001b[1;32m    200\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    203\u001b[0m \u001b[38;5;124;03mthe args/kwargs using `parameters_to_positional_and_keyword` directly\u001b[39;00m\n\u001b[1;32m    204\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    205\u001b[0m args, kwargs \u001b[38;5;241m=\u001b[39m parameters_to_args_kwargs(fn, parameters)\n\u001b[0;32m--> 206\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m fn(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/mlip-arena/mlip_arena/tasks/md.py:363\u001b[0m, in \u001b[0;36mrun\u001b[0;34m()\u001b[0m\n\u001b[1;32m    360\u001b[0m     md_runner\u001b[38;5;241m.\u001b[39mattach(_callback, interval\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m)\n\u001b[1;32m    362\u001b[0m     start_time \u001b[38;5;241m=\u001b[39m datetime\u001b[38;5;241m.\u001b[39mnow()\n\u001b[0;32m--> 363\u001b[0m     md_runner\u001b[38;5;241m.\u001b[39mrun(steps\u001b[38;5;241m=\u001b[39mn_steps)\n\u001b[1;32m    364\u001b[0m     end_time \u001b[38;5;241m=\u001b[39m datetime\u001b[38;5;241m.\u001b[39mnow()\n\u001b[1;32m    366\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m traj_file \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ase/md/npt.py:277\u001b[0m, in \u001b[0;36mrun\u001b[0;34m()\u001b[0m\n\u001b[1;32m    275\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Perform a number of time steps.\"\"\"\u001b[39;00m\n\u001b[1;32m    276\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39minitialized:\n\u001b[0;32m--> 277\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39minitialize()\n\u001b[1;32m    278\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    279\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhave_the_atoms_been_changed():\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ase/md/npt.py:399\u001b[0m, in \u001b[0;36minitialize\u001b[0;34m()\u001b[0m\n\u001b[1;32m    396\u001b[0m deltazeta \u001b[38;5;241m=\u001b[39m dt \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtfact \u001b[38;5;241m*\u001b[39m (atoms\u001b[38;5;241m.\u001b[39mget_kinetic_energy() \u001b[38;5;241m-\u001b[39m\n\u001b[1;32m    397\u001b[0m                                \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdesiredEkin)\n\u001b[1;32m    398\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mzeta_past \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mzeta \u001b[38;5;241m-\u001b[39m deltazeta\n\u001b[0;32m--> 399\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_calculate_q_past_and_future()\n\u001b[1;32m    400\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39minitialized \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ase/md/npt.py:608\u001b[0m, in \u001b[0;36m_calculate_q_past_and_future\u001b[0;34m()\u001b[0m\n\u001b[1;32m    606\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;241m2\u001b[39m):\n\u001b[1;32m    607\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mq_past \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mq \u001b[38;5;241m-\u001b[39m dt \u001b[38;5;241m*\u001b[39m np\u001b[38;5;241m.\u001b[39mdot(p \u001b[38;5;241m/\u001b[39m m, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39minv_h)\n\u001b[0;32m--> 608\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_calculate_q_future(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39matoms\u001b[38;5;241m.\u001b[39mget_forces(md\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m))\n\u001b[1;32m    609\u001b[0m     p \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mdot(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mq_future \u001b[38;5;241m-\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mq_past, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mh \u001b[38;5;241m/\u001b[39m (\u001b[38;5;241m2\u001b[39m \u001b[38;5;241m*\u001b[39m dt)) \u001b[38;5;241m*\u001b[39m m\n\u001b[1;32m    610\u001b[0m     e \u001b[38;5;241m=\u001b[39m ekin(p)\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ase/atoms.py:812\u001b[0m, in \u001b[0;36mget_forces\u001b[0;34m()\u001b[0m\n\u001b[1;32m    810\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_calc \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    811\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mAtoms object has no calculator.\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[0;32m--> 812\u001b[0m forces \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_calc\u001b[38;5;241m.\u001b[39mget_forces(\u001b[38;5;28mself\u001b[39m)\n\u001b[1;32m    814\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m apply_constraint:\n\u001b[1;32m    815\u001b[0m     \u001b[38;5;66;03m# We need a special md flag here because for MD we want\u001b[39;00m\n\u001b[1;32m    816\u001b[0m     \u001b[38;5;66;03m# to skip real constraints but include special \"constraints\"\u001b[39;00m\n\u001b[1;32m    817\u001b[0m     \u001b[38;5;66;03m# Like Hookean.\u001b[39;00m\n\u001b[1;32m    818\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m constraint \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconstraints:\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ase/calculators/abc.py:30\u001b[0m, in \u001b[0;36mget_forces\u001b[0;34m()\u001b[0m\n\u001b[1;32m     29\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mget_forces\u001b[39m(\u001b[38;5;28mself\u001b[39m, atoms\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[0;32m---> 30\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mget_property(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mforces\u001b[39m\u001b[38;5;124m'\u001b[39m, atoms)\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ase/calculators/calculator.py:538\u001b[0m, in \u001b[0;36mget_property\u001b[0;34m()\u001b[0m\n\u001b[1;32m    535\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39muse_cache:\n\u001b[1;32m    536\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39matoms \u001b[38;5;241m=\u001b[39m atoms\u001b[38;5;241m.\u001b[39mcopy()\n\u001b[0;32m--> 538\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcalculate(atoms, [name], system_changes)\n\u001b[1;32m    540\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m name \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mresults:\n\u001b[1;32m    541\u001b[0m     \u001b[38;5;66;03m# For some reason the calculator was not able to do what we want,\u001b[39;00m\n\u001b[1;32m    542\u001b[0m     \u001b[38;5;66;03m# and that is OK.\u001b[39;00m\n\u001b[1;32m    543\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m PropertyNotImplementedError(\n\u001b[1;32m    544\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m not present in this \u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mcalculation\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(name)\n\u001b[1;32m    545\u001b[0m     )\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/mlip-arena/mlip_arena/models/externals/chgnet.py:36\u001b[0m, in \u001b[0;36mcalculate\u001b[0;34m()\u001b[0m\n\u001b[1;32m     30\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcalculate\u001b[39m(\n\u001b[1;32m     31\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m     32\u001b[0m     atoms: Atoms \u001b[38;5;241m|\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     33\u001b[0m     properties: \u001b[38;5;28mlist\u001b[39m \u001b[38;5;241m|\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     34\u001b[0m     system_changes: \u001b[38;5;28mlist\u001b[39m \u001b[38;5;241m|\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     35\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m---> 36\u001b[0m     \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39mcalculate(atoms, properties, system_changes)\n\u001b[1;32m     38\u001b[0m     \u001b[38;5;66;03m# for ase.io.write compatibility\u001b[39;00m\n\u001b[1;32m     39\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mresults\u001b[38;5;241m.\u001b[39mpop(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcrystal_fea\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28;01mNone\u001b[39;00m)\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/chgnet/model/dynamics.py:143\u001b[0m, in \u001b[0;36mcalculate\u001b[0;34m()\u001b[0m\n\u001b[1;32m    141\u001b[0m structure \u001b[38;5;241m=\u001b[39m AseAtomsAdaptor\u001b[38;5;241m.\u001b[39mget_structure(atoms)\n\u001b[1;32m    142\u001b[0m graph \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodel\u001b[38;5;241m.\u001b[39mgraph_converter(structure)\n\u001b[0;32m--> 143\u001b[0m model_prediction \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodel\u001b[38;5;241m.\u001b[39mpredict_graph(\n\u001b[1;32m    144\u001b[0m     graph\u001b[38;5;241m.\u001b[39mto(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdevice), task\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mefsm\u001b[39m\u001b[38;5;124m\"\u001b[39m, return_crystal_feas\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[1;32m    145\u001b[0m )\n\u001b[1;32m    147\u001b[0m \u001b[38;5;66;03m# Convert Result\u001b[39;00m\n\u001b[1;32m    148\u001b[0m factor \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodel\u001b[38;5;241m.\u001b[39mis_intensive \u001b[38;5;28;01melse\u001b[39;00m structure\u001b[38;5;241m.\u001b[39mcomposition\u001b[38;5;241m.\u001b[39mnum_atoms\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/chgnet/model/model.py:627\u001b[0m, in \u001b[0;36mpredict_graph\u001b[0;34m()\u001b[0m\n\u001b[1;32m    625\u001b[0m n_steps \u001b[38;5;241m=\u001b[39m math\u001b[38;5;241m.\u001b[39mceil(\u001b[38;5;28mlen\u001b[39m(graphs) \u001b[38;5;241m/\u001b[39m batch_size)\n\u001b[1;32m    626\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m step \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(n_steps):\n\u001b[0;32m--> 627\u001b[0m     prediction \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mforward(\n\u001b[1;32m    628\u001b[0m         [\n\u001b[1;32m    629\u001b[0m             g\u001b[38;5;241m.\u001b[39mto(model_device)\n\u001b[1;32m    630\u001b[0m             \u001b[38;5;28;01mfor\u001b[39;00m g \u001b[38;5;129;01min\u001b[39;00m graphs[batch_size \u001b[38;5;241m*\u001b[39m step : batch_size \u001b[38;5;241m*\u001b[39m (step \u001b[38;5;241m+\u001b[39m \u001b[38;5;241m1\u001b[39m)]\n\u001b[1;32m    631\u001b[0m         ],\n\u001b[1;32m    632\u001b[0m         task\u001b[38;5;241m=\u001b[39mtask,\n\u001b[1;32m    633\u001b[0m         return_site_energies\u001b[38;5;241m=\u001b[39mreturn_site_energies,\n\u001b[1;32m    634\u001b[0m         return_atom_feas\u001b[38;5;241m=\u001b[39mreturn_atom_feas,\n\u001b[1;32m    635\u001b[0m         return_crystal_feas\u001b[38;5;241m=\u001b[39mreturn_crystal_feas,\n\u001b[1;32m    636\u001b[0m     )\n\u001b[1;32m    637\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m key \u001b[38;5;129;01min\u001b[39;00m {\n\u001b[1;32m    638\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124me\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m    639\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    644\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcrystal_fea\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m    645\u001b[0m     } \u001b[38;5;241m&\u001b[39m {\u001b[38;5;241m*\u001b[39mprediction}:\n\u001b[1;32m    646\u001b[0m         \u001b[38;5;28;01mfor\u001b[39;00m idx, tensor \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(prediction[key]):\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/chgnet/model/model.py:359\u001b[0m, in \u001b[0;36mforward\u001b[0;34m()\u001b[0m\n\u001b[1;32m    354\u001b[0m comp_energy \u001b[38;5;241m=\u001b[39m (\n\u001b[1;32m    355\u001b[0m     \u001b[38;5;241m0\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcomposition_model \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcomposition_model(graphs)\n\u001b[1;32m    356\u001b[0m )\n\u001b[1;32m    358\u001b[0m \u001b[38;5;66;03m# Make batched graph\u001b[39;00m\n\u001b[0;32m--> 359\u001b[0m batched_graph \u001b[38;5;241m=\u001b[39m BatchedGraph\u001b[38;5;241m.\u001b[39mfrom_graphs(\n\u001b[1;32m    360\u001b[0m     graphs,\n\u001b[1;32m    361\u001b[0m     bond_basis_expansion\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbond_basis_expansion,\n\u001b[1;32m    362\u001b[0m     angle_basis_expansion\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mangle_basis_expansion,\n\u001b[1;32m    363\u001b[0m     compute_stress\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124ms\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m task,\n\u001b[1;32m    364\u001b[0m )\n\u001b[1;32m    366\u001b[0m \u001b[38;5;66;03m# Pass to model\u001b[39;00m\n\u001b[1;32m    367\u001b[0m prediction \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compute(\n\u001b[1;32m    368\u001b[0m     batched_graph,\n\u001b[1;32m    369\u001b[0m     compute_force\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m task,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    374\u001b[0m     return_crystal_feas\u001b[38;5;241m=\u001b[39mreturn_crystal_feas,\n\u001b[1;32m    375\u001b[0m )\n",
      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/chgnet/model/model.py:822\u001b[0m, in \u001b[0;36mfrom_graphs\u001b[0;34m()\u001b[0m\n\u001b[1;32m    819\u001b[0m \u001b[38;5;66;03m# Bonds\u001b[39;00m\n\u001b[1;32m    820\u001b[0m atom_cart_coords \u001b[38;5;241m=\u001b[39m graph\u001b[38;5;241m.\u001b[39matom_frac_coord \u001b[38;5;241m@\u001b[39m lattice\n\u001b[1;32m    821\u001b[0m bond_basis_ag, bond_basis_bg, bond_vectors \u001b[38;5;241m=\u001b[39m bond_basis_expansion(\n\u001b[0;32m--> 822\u001b[0m     center\u001b[38;5;241m=\u001b[39matom_cart_coords[graph\u001b[38;5;241m.\u001b[39matom_graph[:, \u001b[38;5;241m0\u001b[39m]],\n\u001b[1;32m    823\u001b[0m     neighbor\u001b[38;5;241m=\u001b[39matom_cart_coords[graph\u001b[38;5;241m.\u001b[39matom_graph[:, \u001b[38;5;241m1\u001b[39m]],\n\u001b[1;32m    824\u001b[0m     undirected2directed\u001b[38;5;241m=\u001b[39mgraph\u001b[38;5;241m.\u001b[39mundirected2directed,\n\u001b[1;32m    825\u001b[0m     image\u001b[38;5;241m=\u001b[39mgraph\u001b[38;5;241m.\u001b[39mneighbor_image,\n\u001b[1;32m    826\u001b[0m     lattice\u001b[38;5;241m=\u001b[39mlattice,\n\u001b[1;32m    827\u001b[0m )\n\u001b[1;32m    828\u001b[0m atom_positions\u001b[38;5;241m.\u001b[39mappend(atom_cart_coords)\n\u001b[1;32m    829\u001b[0m bond_bases_ag\u001b[38;5;241m.\u001b[39mappend(bond_basis_ag)\n",
      "\u001b[0;31mIndexError\u001b[0m: too many indices for tensor of dimension 1"
     ]
    }
   ],
   "source": [
    "results = combustion(atoms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlip_arena"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mlip-arena",
   "language": "python",
   "name": "mlip-arena"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}