Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion arc/job/adapters/xtb_adapter.py
Original file line number Diff line number Diff line change
Expand Up @@ -218,7 +218,9 @@ def write_input_file(self) -> None:
directives, block = '', ''
uhf = self.species[0].number_of_radicals or self.multiplicity - 1

if self.job_type in ['opt', 'conf_opt', 'scan']:
if self.job_type in ['opt', 'conf_opt', 'scan'] \
or (self.job_type == 'directed_scan' and self.directed_scan_type is not None
and 'opt' in self.directed_scan_type):
directives += ' --opt'
directives += self.add_accuracy()
if self.constraints and self.job_type != 'scan':
Expand Down
213 changes: 210 additions & 3 deletions arc/plotter.py
Original file line number Diff line number Diff line change
Expand Up @@ -1293,6 +1293,12 @@ def plot_2d_rotor_scan(results: dict,
if len(results['scans']) != 2:
raise InputError(f'results must represent a 2D rotor, got {len(results["scans"])}D')

# Dispatch to sparse plotting for adaptive scans
if is_sparse_2d_scan(results):
_plot_2d_rotor_scan_sparse(results, path=path, label=label, cmap=cmap,
resolution=resolution, original_dihedrals=original_dihedrals)
return

results['directed_scan'] = clean_scan_results(results['directed_scan'])

# phis0 and phis1 correspond to columns and rows in energies, respectively
Expand Down Expand Up @@ -1357,9 +1363,9 @@ def plot_2d_rotor_scan(results: dict,
label = ' for ' + label if label else ''
plt.title(f'2D scan energies (kJ/mol){label}')
min_x = min_y = -180
plt.xlim = (min_x, min_x + 360)
plt.gca().set_xlim(min_x, min_x + 360)
plt.xticks(np.arange(min_x, min_x + 361, step=60))
plt.ylim = (min_y, min_y + 360)
plt.gca().set_ylim(min_y, min_y + 360)
plt.yticks(np.arange(min_y, min_y + 361, step=60))

if mark_lowest_conformations:
Expand All @@ -1379,6 +1385,207 @@ def plot_2d_rotor_scan(results: dict,
plt.close(fig=fig)


def is_sparse_2d_scan(results: dict) -> bool:
"""
Detect whether a 2D scan results dict represents a sparse/adaptive scan.

A scan is considered sparse if the results contain
``sampling_policy == 'adaptive'``.

Args:
results (dict): The results dictionary from a 2D directed scan.

Returns:
bool: ``True`` if the scan is sparse/adaptive.
"""
return results.get('sampling_policy') == 'adaptive'


def extract_sparse_2d_points(results: dict) -> dict:
"""
Extract sampled point coordinates and energies from a sparse 2D scan result.

Args:
results (dict): The results dictionary from a 2D directed scan.

Returns:
dict: A dictionary with keys ``'x'``, ``'y'``, ``'energy'`` (lists of floats for
completed points with non-None energy), plus ``'failed_points'`` and
``'invalid_points'`` (lists of ``[x, y]`` pairs).
"""
xs, ys, energies = [], [], []
for key, entry in results.get('directed_scan', {}).items():
e = entry.get('energy')
if e is not None:
xs.append(float(key[0]))
ys.append(float(key[1]))
energies.append(float(e))
summary = results.get('adaptive_scan_summary', {})
return {
'x': xs,
'y': ys,
'energy': energies,
'failed_points': summary.get('failed_points', []),
'invalid_points': summary.get('invalid_points', []),
}


def interpolate_sparse_2d_scan(points_x: list,
points_y: list,
energies: list,
grid_resolution: float = 2.0,
) -> tuple:
"""
Interpolate sparse 2D scan data onto a dense grid for contour plotting.

Uses ``scipy.interpolate.griddata`` with periodic boundary augmentation
to reduce artifacts at the -180/+180 wrap boundary.

Args:
points_x (list): Sampled dihedral angles for dimension 0 (degrees).
points_y (list): Sampled dihedral angles for dimension 1 (degrees).
energies (list): Energy values at sampled points (kJ/mol).
grid_resolution (float): Spacing of the dense output grid in degrees.

Returns:
tuple: ``(grid_x, grid_y, grid_energies)`` where each is a 2D numpy array
suitable for ``plt.contourf``.
"""
from scipy.interpolate import griddata

px = np.array(points_x, dtype=np.float64)
py = np.array(points_y, dtype=np.float64)
pe = np.array(energies, dtype=np.float64)

# Augment with periodic image points for wrap-around
aug_x, aug_y, aug_e = list(px), list(py), list(pe)
for dx in (-360.0, 0.0, 360.0):
for dy in (-360.0, 0.0, 360.0):
if dx == 0.0 and dy == 0.0:
continue
aug_x.extend(px + dx)
aug_y.extend(py + dy)
aug_e.extend(pe)
aug_x = np.array(aug_x)
aug_y = np.array(aug_y)
aug_e = np.array(aug_e)

# Dense grid from -180 to 180
n_pts = int(360.0 / grid_resolution) + 1
gx = np.linspace(-180.0, 180.0, n_pts)
gy = np.linspace(-180.0, 180.0, n_pts)
grid_x, grid_y = np.meshgrid(gx, gy, indexing='ij')

# Interpolate: try cubic, fall back to linear, then nearest
pts = np.column_stack([aug_x, aug_y])
grid_e = None
for method in ('cubic', 'linear'):
try:
grid_e = griddata(pts, aug_e, (grid_x, grid_y), method=method)
if not np.all(np.isnan(grid_e)):
break
except (ValueError, Exception):
grid_e = None
if grid_e is None or np.all(np.isnan(grid_e)):
grid_e = griddata(pts, aug_e, (grid_x, grid_y), method='nearest')
# Fill any remaining NaN with nearest-neighbor
mask = np.isnan(grid_e)
if mask.any():
grid_nearest = griddata(pts, aug_e, (grid_x, grid_y), method='nearest')
grid_e[mask] = grid_nearest[mask]

return grid_x, grid_y, grid_e


def _plot_2d_rotor_scan_sparse(results: dict,
path: Optional[str] = None,
label: str = '',
cmap: str = 'Blues',
resolution: int = 90,
original_dihedrals: Optional[List[float]] = None,
):
"""
Plot a sparse/adaptive 2D rotor scan using interpolation for contours
and overlaying sampled, failed, and invalid points.

This is called internally by :func:`plot_2d_rotor_scan` when the results
are detected as sparse.

Args:
results (dict): The results dictionary from a 2D directed scan.
path (str, optional): Folder path to save the plot image.
label (str, optional): Species label.
cmap (str, optional): Matplotlib colormap name.
resolution (int, optional): Image DPI.
original_dihedrals (list, optional): Original dihedral angles for marker.
"""
data = extract_sparse_2d_points(results)
xs, ys, energies = data['x'], data['y'], data['energy']

if len(xs) < 3:
logger.warning(f'Not enough sparse points to plot 2D scan ({len(xs)} points)')
return

# Normalize energies to min = 0
e_min = min(energies)
energies_norm = [e - e_min for e in energies]

# Interpolate to dense grid
grid_x, grid_y, grid_e = interpolate_sparse_2d_scan(xs, ys, energies_norm, grid_resolution=2.0)

fig = plt.figure(num=None, figsize=(12, 8), dpi=resolution, facecolor='w', edgecolor='k')

plt.contourf(grid_x, grid_y, grid_e, 20, cmap=cmap)
plt.colorbar()
contours = plt.contour(grid_x, grid_y, grid_e, 4, colors='black')
plt.clabel(contours, inline=True, fontsize=8)

# Overlay sampled points
plt.scatter(xs, ys, c='black', s=12, zorder=5, label='sampled')

# Overlay failed points
failed = data.get('failed_points', [])
if failed:
fx = [p[0] for p in failed]
fy = [p[1] for p in failed]
plt.scatter(fx, fy, c='red', marker='x', s=40, zorder=6, label='failed')

# Overlay invalid points
invalid = data.get('invalid_points', [])
if invalid:
ix = [p[0] for p in invalid]
iy = [p[1] for p in invalid]
plt.scatter(ix, iy, edgecolors='orange', marker='s', facecolors='none',
s=40, zorder=6, label='invalid')

# Mark original dihedral
if original_dihedrals is not None and len(original_dihedrals) >= 2:
plt.plot(original_dihedrals[0], original_dihedrals[1], color='r',
marker='.', markersize=15, linewidth=0, label='original')

plt.xlabel(f'Dihedral 1 for {results["scans"][0]} (degrees)')
plt.ylabel(f'Dihedral 2 for {results["scans"][1]} (degrees)')
label_str = ' for ' + label if label else ''
summary = results.get('adaptive_scan_summary', {})
n_pts = summary.get('completed_count', len(xs))
plt.title(f'2D scan energies (kJ/mol){label_str} [adaptive, {n_pts} pts]')
plt.gca().set_xlim(-180, 180)
plt.xticks(np.arange(-180, 181, step=60))
plt.gca().set_ylim(-180, 180)
plt.yticks(np.arange(-180, 181, step=60))

plt.legend(loc='upper right', fontsize=8)

if path is not None:
fig_name = f'{results["directed_scan_type"]}_{results["scans"]}_adaptive.png'
fig_path = os.path.join(path, fig_name)
plt.savefig(fig_path, dpi=resolution, facecolor='w', edgecolor='w', orientation='portrait',
format='png', transparent=False, bbox_inches=None, pad_inches=0.1, metadata=None)

plt.show()
plt.close(fig=fig)


def plot_2d_scan_bond_dihedral(results: dict,
path: Optional[str] = None,
label: str = '',
Expand Down Expand Up @@ -1486,7 +1693,7 @@ def plot_2d_scan_bond_dihedral(results: dict,
label = ' for ' + label if label else ''
plt.title(f'2D scan energies (kJ/mol){label}')
min_x = -180
plt.xlim = (min_x, min_x + 360)
plt.gca().set_xlim(min_x, min_x + 360)
plt.xticks(np.arange(min_x, min_x + 361, step=60))

if original_dihedrals is not None:
Expand Down
129 changes: 129 additions & 0 deletions arc/plotter_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
import shutil
import unittest

import numpy as np

import arc.plotter as plotter
from arc.common import ARC_PATH, ARC_TESTING_PATH, read_yaml_file, safe_copy_file
from arc.species.converter import str_to_xyz
Expand Down Expand Up @@ -236,5 +238,132 @@ def tearDownClass(cls):
os.remove(file_path)


class TestSparse2DPlotting(unittest.TestCase):
"""
Contains unit tests for sparse 2D rotor scan plotting helpers.
"""

def _make_dense_results(self):
"""Helper: build a small dense 2D result dict (4x4 grid, increment=120)."""
directed_scan = {}
for a0 in [0.0, 120.0, -120.0, 0.0]:
for a1 in [0.0, 120.0, -120.0, 0.0]:
key = (f'{a0:.2f}', f'{a1:.2f}')
if key not in directed_scan:
directed_scan[key] = {
'energy': float(abs(a0) + abs(a1)) / 10.0,
'xyz': {},
'is_isomorphic': True,
'trsh': [],
}
return {
'directed_scan_type': 'brute_force_opt',
'scans': [[1, 2, 3, 4], [5, 6, 7, 8]],
'directed_scan': directed_scan,
}

def _make_sparse_results(self, n_points=20):
"""Helper: build a sparse adaptive 2D result dict."""
import random
random.seed(42)
directed_scan = {}
for _ in range(n_points):
a0 = round(random.uniform(-180, 180), 2)
a1 = round(random.uniform(-180, 180), 2)
key = (f'{a0:.2f}', f'{a1:.2f}')
directed_scan[key] = {
'energy': float(abs(a0) + abs(a1)) / 10.0,
'xyz': {},
'is_isomorphic': True,
'trsh': [],
}
return {
'directed_scan_type': 'brute_force_opt',
'scans': [[1, 2, 3, 4], [5, 6, 7, 8]],
'directed_scan': directed_scan,
'sampling_policy': 'adaptive',
'adaptive_scan_summary': {
'completed_count': n_points,
'failed_count': 2,
'invalid_count': 1,
'stopping_reason': 'max_points_reached',
'failed_points': [[45.0, -90.0], [120.0, 60.0]],
'invalid_points': [[-30.0, 150.0]],
},
}

def test_is_sparse_2d_scan_dense(self):
"""Test that dense results are not detected as sparse."""
results = self._make_dense_results()
self.assertFalse(plotter.is_sparse_2d_scan(results))

def test_is_sparse_2d_scan_adaptive(self):
"""Test that adaptive results are detected as sparse."""
results = self._make_sparse_results()
self.assertTrue(plotter.is_sparse_2d_scan(results))

def test_extract_sparse_2d_points(self):
"""Test extraction of sparse point data."""
results = self._make_sparse_results(15)
data = plotter.extract_sparse_2d_points(results)
self.assertEqual(len(data['x']), len(data['y']))
self.assertEqual(len(data['x']), len(data['energy']))
self.assertGreater(len(data['x']), 0)
self.assertEqual(len(data['failed_points']), 2)
self.assertEqual(len(data['invalid_points']), 1)

def test_extract_sparse_2d_points_dense(self):
"""Test extraction from dense results (no adaptive summary)."""
results = self._make_dense_results()
data = plotter.extract_sparse_2d_points(results)
self.assertGreater(len(data['x']), 0)
self.assertEqual(data['failed_points'], [])
self.assertEqual(data['invalid_points'], [])

def test_interpolate_sparse_2d_scan(self):
"""Test interpolation produces a dense grid."""
xs = [0.0, 90.0, -90.0, 180.0, -180.0, 45.0, -45.0]
ys = [0.0, 90.0, -90.0, 180.0, -180.0, 45.0, -45.0]
es = [0.0, 5.0, 5.0, 10.0, 10.0, 3.0, 3.0]
gx, gy, ge = plotter.interpolate_sparse_2d_scan(xs, ys, es, grid_resolution=10.0)
# Check shapes match
self.assertEqual(gx.shape, gy.shape)
self.assertEqual(gx.shape, ge.shape)
n = int(360.0 / 10.0) + 1
self.assertEqual(gx.shape, (n, n))
# No NaN values
self.assertFalse(np.any(np.isnan(ge)))

def test_plot_sparse_2d_no_crash(self):
"""Test that plotting a sparse scan doesn't crash."""
import tempfile
results = self._make_sparse_results(30)
with tempfile.TemporaryDirectory() as tmpdir:
# Should not raise
plotter.plot_2d_rotor_scan(results, path=tmpdir)
# Check that a file was created
files = os.listdir(tmpdir)
self.assertTrue(any('adaptive' in f for f in files),
f'Expected adaptive plot file, got: {files}')

def test_plot_dense_2d_unchanged(self):
"""Test that plotting a dense scan still works through the legacy path."""
# This exercises the existing code path; if it crashes, the dense path is broken
results = self._make_dense_results()
# Don't save to disk, just ensure no crash
try:
plotter.plot_2d_rotor_scan(results, path=None)
except (ValueError, KeyError):
# Dense path might fail on this small test grid due to missing points,
# but it should NOT dispatch to sparse path
self.assertFalse(plotter.is_sparse_2d_scan(results))

def test_plot_sparse_too_few_points_no_crash(self):
"""Test that sparse plotting with < 3 points doesn't crash."""
results = self._make_sparse_results(2)
# Should not raise, just warn
plotter.plot_2d_rotor_scan(results, path=None)


if __name__ == '__main__':
unittest.main(testRunner=unittest.TextTestRunner(verbosity=2))
Loading
Loading