From 1e6858e7c4fc387c4e0d7cd80850d770b9746212 Mon Sep 17 00:00:00 2001
From: bartandrews <bartandrews123@gmail.com>
Date: Tue, 12 Nov 2024 17:12:08 +0100
Subject: [PATCH] remove MolecularChain class

---
 python/ffsim/tenpy/__init__.py                |  2 --
 python/ffsim/tenpy/circuits/gates.py          |  2 +-
 python/ffsim/tenpy/hamiltonians/__init__.py   |  2 --
 python/ffsim/tenpy/hamiltonians/lattices.py   | 31 -------------------
 .../hamiltonians/molecular_hamiltonian.py     | 14 +++++++--
 5 files changed, 13 insertions(+), 38 deletions(-)
 delete mode 100644 python/ffsim/tenpy/hamiltonians/lattices.py

diff --git a/python/ffsim/tenpy/__init__.py b/python/ffsim/tenpy/__init__.py
index bb0f692cf..4ab39946b 100644
--- a/python/ffsim/tenpy/__init__.py
+++ b/python/ffsim/tenpy/__init__.py
@@ -22,7 +22,6 @@
     sym_cons_basis,
 )
 from ffsim.tenpy.circuits.lucj_circuit import lucj_circuit_as_mps
-from ffsim.tenpy.hamiltonians.lattices import MolecularChain
 from ffsim.tenpy.hamiltonians.molecular_hamiltonian import MolecularHamiltonianMPOModel
 from ffsim.tenpy.util import product_state_as_mps
 
@@ -33,7 +32,6 @@
     "apply_orbital_rotation",
     "givens_rotation",
     "lucj_circuit_as_mps",
-    "MolecularChain",
     "MolecularHamiltonianMPOModel",
     "num_interaction",
     "num_num_interaction",
diff --git a/python/ffsim/tenpy/circuits/gates.py b/python/ffsim/tenpy/circuits/gates.py
index b57a686d6..c6c62da8d 100644
--- a/python/ffsim/tenpy/circuits/gates.py
+++ b/python/ffsim/tenpy/circuits/gates.py
@@ -428,7 +428,7 @@ def apply_diag_coulomb_evolution(
     """
 
     # extract norb
-    norb, _ = mat.shape
+    _, norb, _ = mat.shape
 
     # unpack alpha-alpha and alpha-beta matrices
     mat_aa, mat_ab = mat
diff --git a/python/ffsim/tenpy/hamiltonians/__init__.py b/python/ffsim/tenpy/hamiltonians/__init__.py
index 68bab34ba..d734edb66 100644
--- a/python/ffsim/tenpy/hamiltonians/__init__.py
+++ b/python/ffsim/tenpy/hamiltonians/__init__.py
@@ -10,10 +10,8 @@
 
 """Classes for converting Hamiltonians to TeNPy MPOModel objects."""
 
-from ffsim.tenpy.hamiltonians.lattices import MolecularChain
 from ffsim.tenpy.hamiltonians.molecular_hamiltonian import MolecularHamiltonianMPOModel
 
 __all__ = [
-    "MolecularChain",
     "MolecularHamiltonianMPOModel",
 ]
diff --git a/python/ffsim/tenpy/hamiltonians/lattices.py b/python/ffsim/tenpy/hamiltonians/lattices.py
deleted file mode 100644
index 54059cba4..000000000
--- a/python/ffsim/tenpy/hamiltonians/lattices.py
+++ /dev/null
@@ -1,31 +0,0 @@
-# (C) Copyright IBM 2024.
-#
-# This code is licensed under the Apache License, Version 2.0. You may
-# obtain a copy of this license in the LICENSE.txt file in the root directory
-# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
-#
-# Any modifications or derivative works of this code must retain this
-# copyright notice, and modified files need to carry a notice indicating
-# that they have been altered from the originals.
-
-import numpy as np
-from tenpy.models.lattice import Lattice
-
-# ignore lowercase argument checks to maintain TeNPy naming conventions
-# ruff: noqa: N803
-
-
-class MolecularChain(Lattice):
-    """Molecular chain."""
-
-    def __init__(self, L, norb, site_a, **kwargs):
-        basis = np.array(([norb, 0.0], [0, 1]))
-        pos = np.array([[i, 0] for i in range(norb)])
-
-        kwargs.setdefault("order", "default")
-        kwargs.setdefault("bc", "open")
-        kwargs.setdefault("bc_MPS", "finite")
-        kwargs.setdefault("basis", basis)
-        kwargs.setdefault("positions", pos)
-
-        super().__init__([L, 1], [site_a] * norb, **kwargs)
diff --git a/python/ffsim/tenpy/hamiltonians/molecular_hamiltonian.py b/python/ffsim/tenpy/hamiltonians/molecular_hamiltonian.py
index 1215b24c8..14714c1ad 100644
--- a/python/ffsim/tenpy/hamiltonians/molecular_hamiltonian.py
+++ b/python/ffsim/tenpy/hamiltonians/molecular_hamiltonian.py
@@ -11,11 +11,11 @@
 from __future__ import annotations
 
 import numpy as np
+from tenpy.models.lattice import Lattice
 from tenpy.models.model import CouplingMPOModel
 from tenpy.networks.site import SpinHalfFermionSite
 
 from ffsim.hamiltonians.molecular_hamiltonian import MolecularHamiltonian
-from ffsim.tenpy.hamiltonians.lattices import MolecularChain
 
 # ignore lowercase variable checks to maintain TeNPy naming conventions
 # ruff: noqa: N806
@@ -34,7 +34,17 @@ def init_lattice(self, params):
         L = params.get("L", 1)
         norb = params.get("norb", 4)
         site = self.init_sites(params)
-        lat = MolecularChain(L, norb, site, basis=[[norb, 0], [0, 1]])
+        basis = np.array(([norb, 0.0], [0, 1]))
+        pos = np.array([[i, 0] for i in range(norb)])
+        lat = Lattice(
+            [L, 1],
+            [site] * norb,
+            order="default",
+            bc="open",
+            bc_MPS="finite",
+            basis=basis,
+            positions=pos,
+        )
         return lat
 
     def init_terms(self, params):