From 371b931c851c614e387713a5be5a6b85042e6848 Mon Sep 17 00:00:00 2001
From: Edward Linscott <edward.linscott@psi.ch>
Date: Mon, 6 Nov 2023 15:05:48 +0100
Subject: [PATCH] Tweaks to allow for 1D and 2D MT corrections

---
 src/koopmans/utils/_io.py                |  2 +-
 src/koopmans/workflows/_koopmans_dscf.py |  5 -----
 src/koopmans/workflows/_workflow.py      | 14 ++++++++++++++
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/src/koopmans/utils/_io.py b/src/koopmans/utils/_io.py
index 557ab5a6f..79c88b485 100644
--- a/src/koopmans/utils/_io.py
+++ b/src/koopmans/utils/_io.py
@@ -52,7 +52,7 @@ def construct_cell_parameters_block(atoms: Atoms) -> Dict[str, Any]:
         params = dict(**cell_to_parameters(atoms.cell))
     else:
         params = {'vectors': [list(row) for row in atoms.cell[:]], 'units': 'angstrom'}
-    params['periodic'] = all(atoms.pbc)
+    params['periodic'] = atoms.pbc
     return params
 
 
diff --git a/src/koopmans/workflows/_koopmans_dscf.py b/src/koopmans/workflows/_koopmans_dscf.py
index da01bf5fd..17da1d08d 100644
--- a/src/koopmans/workflows/_koopmans_dscf.py
+++ b/src/koopmans/workflows/_koopmans_dscf.py
@@ -1034,11 +1034,6 @@ def new_kcp_calculator(self, calc_presets: str = 'dft_init',
         if 'print' in calc.prefix:
             calc.parameters.print_wfc_anion = True
 
-        if self.parameters.mt_correction:
-            calc.parameters.which_compensation = 'tcc'
-        else:
-            calc.parameters.which_compensation = 'none'
-
         # If we are using frozen orbitals, we override the above logic and freeze the variational orbitals
         # post-initialization
         if self.parameters.frozen_orbitals and 'init' not in calc.prefix and not any([s == calc.prefix for s in
diff --git a/src/koopmans/workflows/_workflow.py b/src/koopmans/workflows/_workflow.py
index 0dc31a67d..0a7d0d8e2 100644
--- a/src/koopmans/workflows/_workflow.py
+++ b/src/koopmans/workflows/_workflow.py
@@ -677,6 +677,20 @@ def new_calculator(self,
         # Create the calculator
         calc = calc_class(atoms=copy.deepcopy(self.atoms), **all_kwargs)
 
+        # Add Martyna-Tuckerman settings for kcp calculators
+        if calc.parameters.is_valid('which_compensation'):
+            if self.parameters.mt_correction:
+                if all(self.atoms.pbc):
+                    calc.parameters.which_compensation = 'tcc'
+                elif self.atoms.pbc == [False, False, True]:
+                    calc.parameters.which_compensation = 'tcc1d'
+                elif self.atoms.pbc == [True, True, False]:
+                    calc.parameters.which_compensation = 'tcc2d'
+                else:
+                    raise ValueError('Martyna-Tuckerman correction not implemented for this geometry; for 1D and 2D systems please use z as the unique axis')
+            else:
+                calc.parameters.which_compensation = 'none'
+
         # Add the directory if provided
         if directory is not None:
             calc.directory = directory