From 8708ad34dc126f9c5ac91bffd188f57854f04b07 Mon Sep 17 00:00:00 2001
From: Toon Verstraelen <Toon.Verstraelen@UGent.be>
Date: Thu, 11 Jul 2024 11:24:35 +0200
Subject: [PATCH 1/2] Add prepare_segmented function

---
 docs/example_scripts/convert_fchk_molden.py |  2 +-
 iodata/basis.py                             | 10 ---
 iodata/convert.py                           | 34 ++++++++-
 iodata/formats/fchk.py                      | 10 ++-
 iodata/formats/molden.py                    | 42 ++++++-----
 iodata/formats/molekel.py                   | 22 +++---
 iodata/formats/wfn.py                       | 19 ++---
 iodata/formats/wfx.py                       | 19 ++---
 iodata/overlap.py                           |  6 +-
 iodata/prepare.py                           | 67 ++++++++++++++++-
 iodata/test/common.py                       | 38 +++++++++-
 iodata/test/test_basis.py                   | 45 -----------
 iodata/test/test_convert.py                 | 82 +++++++++++++++++++++
 iodata/test/test_fchk.py                    | 23 ++++--
 iodata/test/test_molden.py                  | 54 +++++++++-----
 iodata/test/test_molekel.py                 | 37 ++++++----
 iodata/test/test_prepare.py                 | 63 +++++++++++++++-
 iodata/test/test_wfn.py                     | 20 +++--
 iodata/test/test_wfx.py                     | 23 ++++--
 19 files changed, 446 insertions(+), 170 deletions(-)

diff --git a/docs/example_scripts/convert_fchk_molden.py b/docs/example_scripts/convert_fchk_molden.py
index 3d03c73f..58bde18e 100755
--- a/docs/example_scripts/convert_fchk_molden.py
+++ b/docs/example_scripts/convert_fchk_molden.py
@@ -5,4 +5,4 @@
 mol = load_one("water.fchk")
 # Here you may put some code to manipulate mol before writing it the data
 # to a different file.
-dump_one(mol, "water.molden")
+dump_one(mol, "water.molden", allow_changes=True)
diff --git a/iodata/basis.py b/iodata/basis.py
index 0cf4a68e..4c7b4317 100644
--- a/iodata/basis.py
+++ b/iodata/basis.py
@@ -239,13 +239,3 @@ class MolecularBasis:
     def nbasis(self) -> int:
         """Number of basis functions."""
         return sum(shell.nbasis for shell in self.shells)
-
-    def get_segmented(self):
-        """Unroll generalized contractions."""
-        shells = []
-        for shell in self.shells:
-            for angmom, kind, coeffs in zip(shell.angmoms, shell.kinds, shell.coeffs.T):
-                shells.append(
-                    Shell(shell.icenter, [angmom], [kind], shell.exponents, coeffs.reshape(-1, 1))
-                )
-        return attrs.evolve(self, shells=shells)
diff --git a/iodata/convert.py b/iodata/convert.py
index 8fb1cf0f..c9d4393c 100644
--- a/iodata/convert.py
+++ b/iodata/convert.py
@@ -23,18 +23,20 @@
 Most of them are used by the prepare module.
 """
 
+import attrs
 import numpy as np
 from numpy.typing import NDArray
 
-from .basis import MolecularBasis
+from .basis import MolecularBasis, Shell
 from .orbitals import MolecularOrbitals
 
 __all__ = (
-    "convert_to_unrestricted",
     "convert_conventions",
     "iter_cart_alphabet",
     "HORTON2_CONVENTIONS",
     "CCA_CONVENTIONS",
+    "convert_to_unrestricted",
+    "convert_to_segmented",
 )
 
 
@@ -253,3 +255,31 @@ def convert_to_unrestricted(mo: MolecularOrbitals) -> MolecularOrbitals:
         None if mo.energies is None else np.concatenate([mo.energies, mo.energies]),
         None if mo.irreps is None else np.concatenate([mo.irreps, mo.irreps]),
     )
+
+
+def convert_to_segmented(obasis: MolecularBasis, keep_sp: bool = False) -> MolecularBasis:
+    """Convert basis with generalized contractions to one with only single contractions.
+
+    Parameters
+    ----------
+    obasis
+        The basis set to convert.
+    keep_sp
+        If True, SP shells are not split up.
+        This can be useful for file formats only support
+        SP-generalized contractions and no other ones.
+
+    Returns
+    -------
+    A new ``MolecularBasis`` instance with separate single contractions.
+    """
+    shells = []
+    for shell in obasis.shells:
+        if (shell.ncon == 1) or (keep_sp and shell.ncon == 2 and (shell.angmoms == [0, 1]).all()):
+            shells.append(shell)
+        else:
+            for angmom, kind, coeffs in zip(shell.angmoms, shell.kinds, shell.coeffs.T):
+                shells.append(
+                    Shell(shell.icenter, [angmom], [kind], shell.exponents, coeffs.reshape(-1, 1))
+                )
+    return attrs.evolve(obasis, shells=shells)
diff --git a/iodata/formats/fchk.py b/iodata/formats/fchk.py
index f1c54902..d34b604f 100644
--- a/iodata/formats/fchk.py
+++ b/iodata/formats/fchk.py
@@ -31,7 +31,8 @@
 from ..docstrings import document_dump_one, document_load_many, document_load_one
 from ..iodata import IOData
 from ..orbitals import MolecularOrbitals
-from ..utils import DumpError, LineIterator, LoadError, LoadWarning, PrepareDumpError, amu
+from ..prepare import prepare_segmented
+from ..utils import LineIterator, LoadError, LoadWarning, PrepareDumpError, amu
 
 __all__ = ()
 
@@ -592,7 +593,7 @@ def prepare_dump(data: IOData, allow_changes: bool, filename: str) -> IOData:
                 "followed by fully virtual ones.",
                 filename,
             )
-    return data
+    return prepare_segmented(data, True, allow_changes, filename, "FCHK")
 
 
 @document_dump_one(
@@ -671,7 +672,10 @@ def dump_one(f: TextIO, data: IOData):
             elif shell.ncon == 2 and (shell.angmoms == [0, 1]).all():
                 shell_types.append(-1)
             else:
-                raise DumpError("Cannot identify type of shell!", f)
+                raise RuntimeError(
+                    "Generalized contractions other than SP are not supported. "
+                    "Call prepare_dump first."
+                )
 
         num_pure_d_shells = sum([1 for st in shell_types if st == 2])
         num_pure_f_shells = sum([1 for st in shell_types if st == 3])
diff --git a/iodata/formats/molden.py b/iodata/formats/molden.py
index 7cd871a1..6d66e9fa 100644
--- a/iodata/formats/molden.py
+++ b/iodata/formats/molden.py
@@ -40,7 +40,7 @@
 from ..orbitals import MolecularOrbitals
 from ..overlap import compute_overlap, gob_cart_normalization
 from ..periodic import num2sym, sym2num
-from ..prepare import prepare_unrestricted_aminusb
+from ..prepare import prepare_segmented, prepare_unrestricted_aminusb
 from ..utils import DumpError, LineIterator, LoadError, LoadWarning, PrepareDumpError, angstrom
 
 __all__ = ()
@@ -793,7 +793,8 @@ def prepare_dump(data: IOData, allow_changes: bool, filename: str) -> IOData:
         raise PrepareDumpError("The Molden format requires an orbital basis set.", filename)
     if data.mo.kind == "generalized":
         raise PrepareDumpError("Cannot write Molden file with generalized orbitals.", filename)
-    return prepare_unrestricted_aminusb(data, allow_changes, filename, "Molden")
+    data = prepare_unrestricted_aminusb(data, allow_changes, filename, "Molden")
+    return prepare_segmented(data, False, allow_changes, filename, "Molden")
 
 
 @document_dump_one("Molden", ["atcoords", "atnums", "mo", "obasis"], ["atcorenums", "title"])
@@ -826,16 +827,19 @@ def dump_one(f: TextIO, data: IOData):
     # to be relevant. If it happens, an error is raised.
     angmom_kinds = {}
     for shell in obasis.shells:
-        for angmom, kind in zip(shell.angmoms, shell.kinds):
-            if angmom in angmom_kinds:
-                if kind != angmom_kinds[angmom]:
-                    raise DumpError(
-                        "Molden format does not support mixed pure+Cartesian functions for one "
-                        "angular momentum.",
-                        f,
-                    )
-            else:
-                angmom_kinds[angmom] = kind
+        if shell.ncon != 1:
+            raise RuntimeError("Generalized contractions not supported. Call prepare_dump first.")
+        kind = shell.kinds[0]
+        angmom = shell.angmoms[0]
+        if angmom in angmom_kinds:
+            if kind != angmom_kinds[angmom]:
+                raise DumpError(
+                    "Molden format does not support mixed pure+Cartesian functions for one "
+                    "angular momentum.",
+                    f,
+                )
+        else:
+            angmom_kinds[angmom] = kind
 
     # Fill in some defaults (Cartesian) for angmom kinds if needed.
     angmom_kinds.setdefault(2, "c")
@@ -863,12 +867,12 @@ def dump_one(f: TextIO, data: IOData):
                 f.write("\n")
             last_icenter = shell.icenter
             f.write("%3i 0\n" % (shell.icenter + 1))
-        # Write out as a segmented basis. Molden format does not support
-        # generalized contractions.
-        for iangmom, angmom in enumerate(shell.angmoms):
-            f.write(f" {angmom_its(angmom):1s}  {shell.nexp:3d} 1.00\n")
-            for exponent, coeff in zip(shell.exponents, shell.coeffs[:, iangmom]):
-                f.write(f"{exponent:20.10f} {coeff:20.10f}\n")
+        # Write out the basis.
+        # It is guaranteed to be segmented when reaching this part of the code.
+        angmom = shell.angmoms[0]
+        f.write(f" {angmom_its(angmom):1s}  {shell.nexp:3d} 1.00\n")
+        for exponent, coeff in zip(shell.exponents, shell.coeffs[:, 0]):
+            f.write(f"{exponent:20.10f} {coeff:20.10f}\n")
     f.write("\n")
 
     # Get the permutation to convert the orbital coefficients to Molden conventions.
@@ -913,7 +917,7 @@ def dump_one(f: TextIO, data: IOData):
             irreps,
         )
     else:
-        raise RuntimeError("This should not happen because of prepare_dump")
+        raise RuntimeError("Generalized orbitals are not support. Call prepare_dump first.")
 
 
 def _dump_helper_orb(f, spin, occs, coeffs, energies, irreps):
diff --git a/iodata/formats/molekel.py b/iodata/formats/molekel.py
index 7bbb915e..baf82052 100644
--- a/iodata/formats/molekel.py
+++ b/iodata/formats/molekel.py
@@ -34,7 +34,7 @@
 from ..docstrings import document_dump_one, document_load_one
 from ..iodata import IOData
 from ..orbitals import MolecularOrbitals
-from ..prepare import prepare_unrestricted_aminusb
+from ..prepare import prepare_segmented, prepare_unrestricted_aminusb
 from ..utils import DumpError, LineIterator, LoadError, LoadWarning, PrepareDumpError, angstrom
 from .molden import CONVENTIONS, _fix_molden_from_buggy_codes
 
@@ -308,7 +308,8 @@ def prepare_dump(data: IOData, allow_changes: bool, filename: str) -> IOData:
         raise PrepareDumpError("The Molekel format requires an orbital basis set.", filename)
     if data.mo.kind == "generalized":
         raise PrepareDumpError("Cannot write Molekel file with generalized orbitals.", filename)
-    return prepare_unrestricted_aminusb(data, allow_changes, filename, "Molekel")
+    data = prepare_unrestricted_aminusb(data, allow_changes, filename, "Molekel")
+    return prepare_segmented(data, False, allow_changes, filename, "Molekel")
 
 
 @document_dump_one("Molekel", ["atcoords", "atnums", "mo", "obasis"], ["atcharges"])
@@ -346,15 +347,18 @@ def dump_one(f: TextIO, data: IOData):
     f.write("$BASIS\n")
     iatom_last = 0
     for shell in data.obasis.shells:
+        if shell.ncon != 1:
+            raise RuntimeError("Generalized contractions not supported. Call prepare_dump first.")
         iatom_new = shell.icenter
         if iatom_new != iatom_last:
             f.write("$$\n")
-        for iangmom, (angmom, kind) in enumerate(zip(shell.angmoms, shell.kinds)):
-            iatom_last = shell.icenter
-            nbasis = len(CONVENTIONS[(angmom, kind)])
-            f.write(f" {nbasis} {angmom_its(angmom).capitalize():1s} 1.00\n")
-            for exponent, coeff in zip(shell.exponents, shell.coeffs[:, iangmom]):
-                f.write(f"{exponent:20.10f} {coeff:17.10f}\n")
+        angmom = shell.angmoms[0]
+        kind = shell.kinds[0]
+        iatom_last = shell.icenter
+        nbasis = len(CONVENTIONS[(angmom, kind)])
+        f.write(f" {nbasis} {angmom_its(angmom).capitalize():1s} 1.00\n")
+        for exponent, coeff in zip(shell.exponents, shell.coeffs[:, 0]):
+            f.write(f"{exponent:20.10f} {coeff:17.10f}\n")
     f.write("\n")
     f.write("$END\n")
     f.write("\n")
@@ -388,7 +392,7 @@ def dump_one(f: TextIO, data: IOData):
         _dump_helper_occ(f, data, spin="b")
 
     else:
-        raise RuntimeError("This should not happen because of prepare_dump")
+        raise RuntimeError("Generalized orbitals are not support. Call prepare_dump first.")
 
 
 # Defining help dumping functions
diff --git a/iodata/formats/wfn.py b/iodata/formats/wfn.py
index ced8fadb..d84bcf05 100644
--- a/iodata/formats/wfn.py
+++ b/iodata/formats/wfn.py
@@ -39,7 +39,7 @@
 from ..orbitals import MolecularOrbitals
 from ..overlap import gob_cart_normalization
 from ..periodic import num2sym, sym2num
-from ..prepare import prepare_unrestricted_aminusb
+from ..prepare import prepare_segmented, prepare_unrestricted_aminusb
 from ..utils import LineIterator, LoadError, PrepareDumpError
 
 __all__ = ()
@@ -537,7 +537,8 @@ def prepare_dump(data: IOData, allow_changes: bool, filename: str) -> IOData:
             raise PrepareDumpError(
                 "The WFN format only supports Cartesian MolecularBasis.", filename
             )
-    return prepare_unrestricted_aminusb(data, allow_changes, filename, "WFN")
+    data = prepare_unrestricted_aminusb(data, allow_changes, filename, "WFN")
+    return prepare_segmented(data, False, allow_changes, filename, "WFN")
 
 
 @document_dump_one(
@@ -550,13 +551,13 @@ def dump_one(f: TextIO, data: IOData) -> None:
     # get shells for the de-contracted basis
     shells = []
     for shell in data.obasis.shells:
-        for i, (angmom, kind) in enumerate(zip(shell.angmoms, shell.kinds)):
-            for exponent, coeff in zip(shell.exponents, shell.coeffs.T[i]):
-                shells.append(
-                    Shell(
-                        shell.icenter, [angmom], [kind], np.array([exponent]), coeff.reshape(-1, 1)
-                    )
-                )
+        if shell.ncon != 1:
+            raise RuntimeError("Generalized contractions not supported. Call prepare_dump first.")
+        # Decontract the shell
+        angmom = shell.angmoms[0]
+        kind = shell.kinds[0]
+        for exponent, coeff in zip(shell.exponents, shell.coeffs[:, 0]):
+            shells.append(Shell(shell.icenter, [angmom], [kind], [exponent], [[coeff]]))
     # make a new instance of MolecularBasis with de-contracted basis shells; ideally for WFN we
     # want the primitive basis set, but IOData only supports shells.
     obasis = MolecularBasis(shells, data.obasis.conventions, data.obasis.primitive_normalization)
diff --git a/iodata/formats/wfx.py b/iodata/formats/wfx.py
index 73c7e18b..e5789035 100644
--- a/iodata/formats/wfx.py
+++ b/iodata/formats/wfx.py
@@ -33,7 +33,7 @@
 from ..iodata import IOData
 from ..orbitals import MolecularOrbitals
 from ..periodic import num2sym
-from ..prepare import prepare_unrestricted_aminusb
+from ..prepare import prepare_segmented, prepare_unrestricted_aminusb
 from ..utils import LineIterator, LoadError, LoadWarning, PrepareDumpError
 from .wfn import CONVENTIONS, build_obasis, get_mocoeff_scales
 
@@ -369,7 +369,8 @@ def prepare_dump(data: IOData, allow_changes: bool, filename: str) -> IOData:
             raise PrepareDumpError(
                 "The WFX format only supports Cartesian MolecularBasis.", filename
             )
-    return prepare_unrestricted_aminusb(data, allow_changes, filename, "WFX")
+    data = prepare_unrestricted_aminusb(data, allow_changes, filename, "WFX")
+    return prepare_segmented(data, False, allow_changes, filename, "WFX")
 
 
 @document_dump_one(
@@ -390,13 +391,13 @@ def dump_one(f: TextIO, data: IOData):
     # get shells for the de-contracted basis
     shells = []
     for shell in data.obasis.shells:
-        for i, (angmom, kind) in enumerate(zip(shell.angmoms, shell.kinds)):
-            for exponent, coeff in zip(shell.exponents, shell.coeffs.T[i]):
-                shells.append(
-                    Shell(
-                        shell.icenter, [angmom], [kind], np.array([exponent]), coeff.reshape(-1, 1)
-                    )
-                )
+        if shell.ncon != 1:
+            raise RuntimeError("Generalized contractions not supported. Call prepare_dump first.")
+        # Decontract the shell
+        angmom = shell.angmoms[0]
+        kind = shell.kinds[0]
+        for exponent, coeff in zip(shell.exponents, shell.coeffs[:, 0]):
+            shells.append(Shell(shell.icenter, [angmom], [kind], [exponent], [[coeff]]))
     # make a new instance of MolecularBasis with de-contracted basis shells; ideally for WFX we
     # want the primitive basis set, but IOData only supports shells.
     obasis = MolecularBasis(shells, data.obasis.conventions, data.obasis.primitive_normalization)
diff --git a/iodata/overlap.py b/iodata/overlap.py
index f62c62c4..dd784ad5 100644
--- a/iodata/overlap.py
+++ b/iodata/overlap.py
@@ -27,7 +27,7 @@
 
 from .basis import MolecularBasis, Shell
 from .convert import HORTON2_CONVENTIONS as OVERLAP_CONVENTIONS
-from .convert import convert_conventions, iter_cart_alphabet
+from .convert import convert_conventions, convert_to_segmented, iter_cart_alphabet
 from .overlap_cartpure import tfs
 
 __all__ = ("OVERLAP_CONVENTIONS", "compute_overlap", "gob_cart_normalization")
@@ -105,7 +105,7 @@ def compute_overlap(
         raise ValueError("The overlap integrals are only implemented for L2 normalization.")
 
     # Get a segmented basis, for simplicity
-    obasis0 = obasis0.get_segmented()
+    obasis0 = convert_to_segmented(obasis0)
 
     # Handle optional arguments
     if obasis1 is None:
@@ -126,7 +126,7 @@ def compute_overlap(
                 "array of atomic coordinates is expected."
             )
         # Get a segmented basis, for simplicity
-        obasis1 = obasis1.get_segmented()
+        obasis1 = convert_to_segmented(obasis1)
         identical = False
 
     # Initialize result
diff --git a/iodata/prepare.py b/iodata/prepare.py
index a95aee73..eafc2d5d 100644
--- a/iodata/prepare.py
+++ b/iodata/prepare.py
@@ -30,11 +30,11 @@
 
 import attrs
 
-from .convert import convert_to_unrestricted
+from .convert import convert_to_segmented, convert_to_unrestricted
 from .iodata import IOData
 from .utils import PrepareDumpError, PrepareDumpWarning
 
-__all__ = ("prepare_unrestricted_aminusb",)
+__all__ = ("prepare_unrestricted_aminusb", "prepare_segmented")
 
 
 def prepare_unrestricted_aminusb(data: IOData, allow_changes: bool, filename: str, fmt: str):
@@ -51,7 +51,6 @@ def prepare_unrestricted_aminusb(data: IOData, allow_changes: bool, filename: st
     fmt
         The file format whose dump function is calling this function, only used for error messages.
 
-
     Returns
     -------
     data
@@ -83,7 +82,7 @@ def prepare_unrestricted_aminusb(data: IOData, allow_changes: bool, filename: st
     message = f"The {fmt} format does not support restricted orbitals with mo.occs_aminusb. "
     if not allow_changes:
         raise PrepareDumpError(
-            message + "Set allow_change=True to enable conversion to unrestricted.", filename
+            message + "Set allow_changes to enable conversion to unrestricted.", filename
         )
     warn(
         PrepareDumpWarning(message + "The orbitals are converted to unrestricted", filename),
@@ -92,3 +91,63 @@ def prepare_unrestricted_aminusb(data: IOData, allow_changes: bool, filename: st
 
     # Convert
     return attrs.evolve(data, mo=convert_to_unrestricted(data.mo))
+
+
+def prepare_segmented(data: IOData, keep_sp: bool, allow_changes: bool, filename: str, fmt: str):
+    """If needed, convert generalized contractions to segmented ones.
+
+    Parameters
+    ----------
+    data
+        The IOData instance with the orbital basis set.
+    keep_sp
+        Set to True if SP-shells should not be segmented.
+    allow_changes
+        Whether conversion of the IOData object to a compatible form is allowed or not.
+    filename
+        The file to be written to, only used for error messages.
+    fmt
+        The file format whose dump function is calling this function, only used for error messages.
+
+    Returns
+    -------
+    data
+        The given data object if no conversion took place,
+        or a shallow copy with some new attriubtes.
+
+    Raises
+    ------
+    ValueError
+        If the given data object has no orbital basis set.
+    PrepareDumpError
+        If ``allow_changes == False`` and a conversion is required.
+    PrepareDumpWarning
+        If ``allow_changes == True`` and a conversion is required.
+    """
+    # Check: possible, needed?
+    if data.obasis is None:
+        raise ValueError("The given IOData instance has no orbital basis set.")
+    if all(
+        shell.ncon == 1 or (keep_sp and shell.ncon == 2 and (shell.angmoms == [0, 1]).all())
+        for shell in data.obasis.shells
+    ):
+        return data
+
+    # Raise error or warning
+    message = f"The {fmt} format does not support generalized contractions"
+    if keep_sp:
+        message += " other than SP shells"
+    message += ". "
+    if not allow_changes:
+        raise PrepareDumpError(
+            message + "Set allow_changes to enable conversion to segmented shells.", filename
+        )
+    warn(
+        PrepareDumpWarning(
+            message + "The orbital basis is converted to segmented shells", filename
+        ),
+        stacklevel=2,
+    )
+
+    # Convert
+    return attrs.evolve(data, obasis=convert_to_segmented(data.obasis, keep_sp))
diff --git a/iodata/test/common.py b/iodata/test/common.py
index 471b766f..4b6e1db2 100644
--- a/iodata/test/common.py
+++ b/iodata/test/common.py
@@ -42,7 +42,8 @@
     "compare_mols",
     "check_orthonormal",
     "load_one_warning",
-    "create_generalized",
+    "create_generalized_orbitals",
+    "create_generalized_contraction",
 )
 
 
@@ -201,7 +202,7 @@ def load_one_warning(
             return load_one(str(fn), fmt=fmt, **kwargs)
 
 
-def create_generalized() -> IOData:
+def create_generalized_orbitals() -> IOData:
     """Create a dummy IOData object with generalized molecular orbitals."""
     rng = np.random.default_rng()
     return IOData(
@@ -213,9 +214,40 @@ def create_generalized() -> IOData:
         obasis=MolecularBasis(
             [
                 Shell(0, [0, 1], ["c", "c"], rng.uniform(0, 1, 2), rng.uniform(0, 1, (2, 2))),
-                Shell(0, [0, 1], ["c", "c"], rng.uniform(0, 1, 2), rng.uniform(0, 1, (2, 2))),
+                Shell(1, [0, 1], ["c", "c"], rng.uniform(0, 1, 2), rng.uniform(0, 1, (2, 2))),
             ],
             {(0, "c"): ["1"], (1, "c"): ["x", "y", "z"]},
             "L2",
         ),
     )
+
+
+def create_generalized_contraction() -> IOData:
+    """Create a dummy IOData object with generalized contractions in the basis."""
+    rng = np.random.default_rng()
+    return IOData(
+        atnums=[1, 1],
+        atcoords=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]],
+        mo=MolecularOrbitals(
+            "restricted",
+            3,
+            3,
+            occs=[1.0, 1.0, 0.0],
+            energies=[0.2, 0.3, 0.4],
+            coeffs=rng.uniform(0, 1, (6, 3)),
+        ),
+        obasis=MolecularBasis(
+            [
+                Shell(
+                    0, [0, 0, 0], ["c", "c", "c"], rng.uniform(0, 1, 4), rng.uniform(0, 1, (4, 3))
+                ),
+                Shell(
+                    1, [0, 0, 0], ["c", "c", "c"], rng.uniform(0, 1, 4), rng.uniform(0, 1, (4, 3))
+                ),
+            ],
+            {
+                (0, "c"): ["1"],
+            },
+            "L2",
+        ),
+    )
diff --git a/iodata/test/test_basis.py b/iodata/test/test_basis.py
index 3a668cbf..0dbd6ebc 100644
--- a/iodata/test/test_basis.py
+++ b/iodata/test/test_basis.py
@@ -21,7 +21,6 @@
 import attrs
 import numpy as np
 import pytest
-from numpy.testing import assert_equal
 
 from ..basis import (
     MolecularBasis,
@@ -163,47 +162,3 @@ def test_nbasis1():
         "L2",
     )
     assert obasis.nbasis == 9
-
-
-def test_get_segmented():
-    rng = np.random.default_rng(1)
-    obasis0 = MolecularBasis(
-        [
-            Shell(0, [0, 1], ["c", "c"], rng.uniform(0, 1, 5), rng.uniform(-1, 1, (5, 2))),
-            Shell(1, [2, 3], ["p", "p"], rng.uniform(0, 1, 7), rng.uniform(-1, 1, (7, 2))),
-        ],
-        CP2K_CONVENTIONS,
-        "L2",
-    )
-    assert obasis0.nbasis == 16
-    obasis1 = obasis0.get_segmented()
-    assert len(obasis1.shells) == 4
-    assert obasis1.nbasis == 16
-    # shell 0
-    shell0 = obasis1.shells[0]
-    assert shell0.icenter == 0
-    assert_equal(shell0.angmoms, [0])
-    assert shell0.kinds == ["c"]
-    assert_equal(shell0.exponents, obasis0.shells[0].exponents)
-    assert_equal(shell0.coeffs, obasis0.shells[0].coeffs[:, :1])
-    # shell 1
-    shell1 = obasis1.shells[1]
-    assert shell1.icenter == 0
-    assert_equal(shell1.angmoms, [1])
-    assert shell1.kinds == ["c"]
-    assert_equal(shell1.exponents, obasis0.shells[0].exponents)
-    assert_equal(shell1.coeffs, obasis0.shells[0].coeffs[:, 1:])
-    # shell 2
-    shell2 = obasis1.shells[2]
-    assert shell2.icenter == 1
-    assert_equal(shell2.angmoms, [2])
-    assert shell2.kinds == ["p"]
-    assert_equal(shell2.exponents, obasis0.shells[1].exponents)
-    assert_equal(shell2.coeffs, obasis0.shells[1].coeffs[:, :1])
-    # shell 0
-    shell3 = obasis1.shells[3]
-    assert shell3.icenter == 1
-    assert_equal(shell3.angmoms, [3])
-    assert shell3.kinds == ["p"]
-    assert_equal(shell3.exponents, obasis0.shells[1].exponents)
-    assert_equal(shell3.coeffs, obasis0.shells[1].coeffs[:, 1:])
diff --git a/iodata/test/test_convert.py b/iodata/test/test_convert.py
index 23f0f43e..2afa2d17 100644
--- a/iodata/test/test_convert.py
+++ b/iodata/test/test_convert.py
@@ -28,9 +28,11 @@
     HORTON2_CONVENTIONS,
     _convert_convention_shell,
     convert_conventions,
+    convert_to_segmented,
     convert_to_unrestricted,
     iter_cart_alphabet,
 )
+from ..formats.cp2klog import CONVENTIONS as CP2K_CONVENTIONS
 from ..orbitals import MolecularOrbitals
 
 
@@ -273,3 +275,83 @@ def test_convert_to_unrestricted_full():
     assert_allclose(mo2.energiesb, mo1.energiesb)
     assert_equal(mo2.irrepsa, mo1.irrepsa)
     assert_equal(mo2.irrepsb, mo1.irrepsb)
+
+
+def test_convert_to_segmented():
+    rng = np.random.default_rng(1)
+    obasis0 = MolecularBasis(
+        [
+            Shell(0, [0, 1], ["c", "c"], rng.uniform(0, 1, 5), rng.uniform(-1, 1, (5, 2))),
+            Shell(1, [2, 3], ["p", "p"], rng.uniform(0, 1, 7), rng.uniform(-1, 1, (7, 2))),
+        ],
+        CP2K_CONVENTIONS,
+        "L2",
+    )
+    assert obasis0.nbasis == 16
+    obasis1 = convert_to_segmented(obasis0)
+    assert len(obasis1.shells) == 4
+    assert obasis1.nbasis == 16
+    # shell 0
+    shell0 = obasis1.shells[0]
+    assert shell0.icenter == 0
+    assert_equal(shell0.angmoms, [0])
+    assert_equal(shell0.kinds, ["c"])
+    assert_equal(shell0.exponents, obasis0.shells[0].exponents)
+    assert_equal(shell0.coeffs, obasis0.shells[0].coeffs[:, :1])
+    # shell 1
+    shell1 = obasis1.shells[1]
+    assert shell1.icenter == 0
+    assert_equal(shell1.angmoms, [1])
+    assert_equal(shell1.kinds, ["c"])
+    assert_equal(shell1.exponents, obasis0.shells[0].exponents)
+    assert_equal(shell1.coeffs, obasis0.shells[0].coeffs[:, 1:])
+    # shell 2
+    shell2 = obasis1.shells[2]
+    assert shell2.icenter == 1
+    assert_equal(shell2.angmoms, [2])
+    assert_equal(shell2.kinds, ["p"])
+    assert_equal(shell2.exponents, obasis0.shells[1].exponents)
+    assert_equal(shell2.coeffs, obasis0.shells[1].coeffs[:, :1])
+    # shell 0
+    shell3 = obasis1.shells[3]
+    assert shell3.icenter == 1
+    assert_equal(shell3.angmoms, [3])
+    assert_equal(shell3.kinds, ["p"])
+    assert_equal(shell3.exponents, obasis0.shells[1].exponents)
+    assert_equal(shell3.coeffs, obasis0.shells[1].coeffs[:, 1:])
+
+
+def test_convert_to_segmented_sp():
+    rng = np.random.default_rng(1)
+    obasis0 = MolecularBasis(
+        [
+            Shell(0, [0, 1], ["c", "c"], rng.uniform(0, 1, 5), rng.uniform(-1, 1, (5, 2))),
+            Shell(1, [2, 3], ["p", "p"], rng.uniform(0, 1, 7), rng.uniform(-1, 1, (7, 2))),
+        ],
+        HORTON2_CONVENTIONS,
+        "L2",
+    )
+    obasis1 = convert_to_segmented(obasis0, keep_sp=True)
+    assert len(obasis1.shells) == 3
+    assert obasis1.nbasis == 16
+    # shell 0
+    shell0 = obasis1.shells[0]
+    assert shell0.icenter == 0
+    assert_equal(shell0.angmoms, [0, 1])
+    assert_equal(shell0.kinds, ["c", "c"])
+    assert_equal(shell0.exponents, obasis0.shells[0].exponents)
+    assert_equal(shell0.coeffs, obasis0.shells[0].coeffs)
+    # shell 1
+    shell2 = obasis1.shells[1]
+    assert shell2.icenter == 1
+    assert_equal(shell2.angmoms, [2])
+    assert_equal(shell2.kinds, ["p"])
+    assert_equal(shell2.exponents, obasis0.shells[1].exponents)
+    assert_equal(shell2.coeffs, obasis0.shells[1].coeffs[:, :1])
+    # shell 3
+    shell3 = obasis1.shells[2]
+    assert shell3.icenter == 1
+    assert_equal(shell3.angmoms, [3])
+    assert_equal(shell3.kinds, ["p"])
+    assert_equal(shell3.exponents, obasis0.shells[1].exponents)
+    assert_equal(shell3.coeffs, obasis0.shells[1].coeffs[:, 1:])
diff --git a/iodata/test/test_fchk.py b/iodata/test/test_fchk.py
index 57aa8592..340c3e63 100644
--- a/iodata/test/test_fchk.py
+++ b/iodata/test/test_fchk.py
@@ -28,12 +28,13 @@
 
 from ..api import dump_one, load_many, load_one
 from ..overlap import compute_overlap
-from ..utils import PrepareDumpError, check_dm
+from ..utils import PrepareDumpError, PrepareDumpWarning, check_dm
 from .common import (
     check_orthonormal,
     compare_mols,
     compute_1rdm,
-    create_generalized,
+    create_generalized_contraction,
+    create_generalized_orbitals,
     load_one_warning,
 )
 from .test_molekel import compare_mols_diff_formats
@@ -698,8 +699,20 @@ def test_dump_fchk_from_molekel(tmpdir, path, match):
     check_load_dump_consistency(tmpdir, path, match)
 
 
-def test_generalized():
+def test_generalized_orbitals():
     # The FCHK format does not support generalized MOs
-    data = create_generalized()
+    data = create_generalized_orbitals()
     with pytest.raises(PrepareDumpError):
-        dump_one(data, "generalized.fchk")
+        dump_one(data, "generalized_orbitals.fchk")
+
+
+def test_fchk_generalized_contraction(tmpdir):
+    data0 = create_generalized_contraction()
+    path_fchk = os.path.join(tmpdir, "generalized_contraction.fchk")
+    with pytest.raises(PrepareDumpError):
+        dump_one(data0, path_fchk)
+    assert not os.path.isfile(path_fchk)
+    with pytest.warns(PrepareDumpWarning):
+        dump_one(data0, path_fchk, allow_changes=True)
+    data1 = load_one(path_fchk)
+    assert all(shell.ncon == 1 for shell in data1.obasis.shells)
diff --git a/iodata/test/test_molden.py b/iodata/test/test_molden.py
index 36a69aaf..69803710 100644
--- a/iodata/test/test_molden.py
+++ b/iodata/test/test_molden.py
@@ -29,14 +29,26 @@
 
 from ..api import dump_one, load_one
 from ..basis import MolecularBasis, Shell
-from ..convert import HORTON2_CONVENTIONS, convert_conventions
+from ..convert import HORTON2_CONVENTIONS, convert_conventions, convert_to_segmented
 from ..formats.molden import _load_low
 from ..formats.molden import dump_one as molden_dump_one
 from ..iodata import IOData
 from ..orbitals import MolecularOrbitals
 from ..overlap import OVERLAP_CONVENTIONS, compute_overlap
-from ..utils import DumpError, LineIterator, LoadWarning, PrepareDumpError, angstrom
-from .common import check_orthonormal, compare_mols, compute_mulliken_charges, create_generalized
+from ..utils import (
+    DumpError,
+    LineIterator,
+    LoadWarning,
+    PrepareDumpError,
+    PrepareDumpWarning,
+    angstrom,
+)
+from .common import (
+    check_orthonormal,
+    compare_mols,
+    compute_mulliken_charges,
+    create_generalized_orbitals,
+)
 
 
 @pytest.mark.slow()
@@ -563,21 +575,21 @@ def test_load_molden_f():
 
 
 @pytest.mark.parametrize(
-    ("fn", "match"),
+    ("fn", "match", "allow_changes"),
     [
-        ("h2o.molden.input", "ORCA"),
-        pytest.param("li2.molden.input", "ORCA", marks=pytest.mark.slow),
-        ("F.molden", "PSI4"),
-        ("nh3_molden_pure.molden", None),
-        ("nh3_molden_cart.molden", None),
-        ("he2_ghost_psi4_1.0.molden", None),
-        pytest.param("psi4_cuh_cc_pvqz_pure.molden", "unnormalized", marks=pytest.mark.slow),
-        ("hf_sto3g.fchk", None),
-        ("h_sto3g.fchk", None),
-        ("ch3_rohf_sto3g_g03.fchk", None),
+        ("h2o.molden.input", "ORCA", False),
+        pytest.param("li2.molden.input", "ORCA", False, marks=pytest.mark.slow),
+        ("F.molden", "PSI4", False),
+        ("nh3_molden_pure.molden", None, False),
+        ("nh3_molden_cart.molden", None, False),
+        ("he2_ghost_psi4_1.0.molden", None, False),
+        pytest.param("psi4_cuh_cc_pvqz_pure.molden", "unnormalized", False, marks=pytest.mark.slow),
+        ("hf_sto3g.fchk", None, True),
+        ("h_sto3g.fchk", None, False),
+        ("ch3_rohf_sto3g_g03.fchk", None, True),
     ],
 )
-def test_load_dump_consistency(tmpdir, fn, match):
+def test_load_dump_consistency(tmpdir, fn, match, allow_changes):
     with as_file(files("iodata.test.data").joinpath(fn)) as file_name:
         if match is None:
             mol1 = load_one(str(file_name))
@@ -585,12 +597,16 @@ def test_load_dump_consistency(tmpdir, fn, match):
             with pytest.warns(LoadWarning, match=match):
                 mol1 = load_one(str(file_name))
     fn_tmp = os.path.join(tmpdir, "foo.bar")
-    dump_one(mol1, fn_tmp, fmt="molden")
+    if allow_changes:
+        with pytest.warns(PrepareDumpWarning):
+            dump_one(mol1, fn_tmp, fmt="molden", allow_changes=True)
+    else:
+        dump_one(mol1, fn_tmp, fmt="molden")
     mol2 = load_one(fn_tmp, fmt="molden")
     # Remove and or fix some things in mol1 to make it compatible with what
     # can be read from a Molden file:
     # - Change basis of mol1 to segmented.
-    mol1.obasis = mol1.obasis.get_segmented()
+    mol1.obasis = convert_to_segmented(mol1.obasis)
     # - Set default irreps in mol1, if not present.
     if mol1.mo.irreps is None:
         mol1.mo = attrs.evolve(mol1.mo, irreps=["1a"] * mol1.mo.norb)
@@ -599,9 +615,9 @@ def test_load_dump_consistency(tmpdir, fn, match):
     compare_mols(mol1, mol2)
 
 
-def test_generalized():
+def test_generalized_orbitals():
     # The Molden format does not support generalized MOs
-    data = create_generalized()
+    data = create_generalized_orbitals()
     with pytest.raises(PrepareDumpError):
         dump_one(data, "generalized.molden")
 
diff --git a/iodata/test/test_molekel.py b/iodata/test/test_molekel.py
index e41ac844..19eeb58c 100644
--- a/iodata/test/test_molekel.py
+++ b/iodata/test/test_molekel.py
@@ -29,12 +29,12 @@
 from ..api import dump_one, load_one
 from ..convert import convert_conventions
 from ..overlap import compute_overlap
-from ..utils import LoadWarning, PrepareDumpError, angstrom
+from ..utils import LoadWarning, PrepareDumpError, PrepareDumpWarning, angstrom
 from .common import (
     check_orthonormal,
     compare_mols,
     compute_mulliken_charges,
-    create_generalized,
+    create_generalized_orbitals,
     load_one_warning,
 )
 
@@ -59,7 +59,9 @@ def compare_mols_diff_formats(mol1, mol2):
     assert_allclose(charges1, charges2, rtol=0.0, atol=1.0e-6)
 
 
-def check_load_dump_consistency(fn: str, tmpdir: str, match: Optional[str] = None):
+def check_load_dump_consistency(
+    fn: str, tmpdir: str, match: Optional[str] = None, allow_changes: bool = False
+):
     """Check if data is preserved after dumping and loading a Molekel file.
 
     Parameters
@@ -71,11 +73,18 @@ def check_load_dump_consistency(fn: str, tmpdir: str, match: Optional[str] = Non
     match
         When given, loading the file is expected to raise a warning whose
         message string contains match.
+    allow_changes
+        Whether to allow changes to the data when writing the file.
+        When True, warnings related to the changes are tested.
 
     """
     mol1 = load_one_warning(fn, match=match)
     fn_tmp = os.path.join(tmpdir, "foo.bar")
-    dump_one(mol1, fn_tmp, fmt="molekel")
+    if allow_changes:
+        with pytest.warns(PrepareDumpWarning):
+            dump_one(mol1, fn_tmp, fmt="molekel", allow_changes=True)
+    else:
+        dump_one(mol1, fn_tmp, fmt="molekel")
     mol2 = load_one(fn_tmp, fmt="molekel")
     form = fn.split(".")
     if "molden" in form or "fchk" in form:
@@ -85,17 +94,17 @@ def check_load_dump_consistency(fn: str, tmpdir: str, match: Optional[str] = Non
 
 
 @pytest.mark.parametrize(
-    ("path", "match"),
+    ("path", "match", "allow_changes"),
     [
-        ("h2_sto3g.mkl", "ORCA"),
-        pytest.param("ethanol.mkl", "ORCA", marks=pytest.mark.slow),
-        pytest.param("li2.mkl", "ORCA", marks=pytest.mark.slow),
-        pytest.param("li2.molden.input", "ORCA", marks=pytest.mark.slow),
-        ("li2_g09_nbasis_indep.fchk", None),
+        ("h2_sto3g.mkl", "ORCA", False),
+        pytest.param("ethanol.mkl", "ORCA", False, marks=pytest.mark.slow),
+        pytest.param("li2.mkl", "ORCA", False, marks=pytest.mark.slow),
+        pytest.param("li2.molden.input", "ORCA", False, marks=pytest.mark.slow),
+        ("li2_g09_nbasis_indep.fchk", None, True),
     ],
 )
-def test_load_dump_consistency(tmpdir, path, match):
-    check_load_dump_consistency(path, tmpdir, match)
+def test_load_dump_consistency(tmpdir, path, match, allow_changes):
+    check_load_dump_consistency(path, tmpdir, match, allow_changes)
 
 
 def test_load_mkl_ethanol():
@@ -162,9 +171,9 @@ def test_load_mkl_h2_huge_threshold():
         load_one(str(fn_molekel), norm_threshold=1e4)
 
 
-def test_generalized():
+def test_generalized_orbitals():
     # The Molden format does not support generalized MOs
-    data = create_generalized()
+    data = create_generalized_orbitals()
     with pytest.raises(PrepareDumpError):
         dump_one(data, "generalized.mkl")
 
diff --git a/iodata/test/test_prepare.py b/iodata/test/test_prepare.py
index dac92aea..be2462ae 100644
--- a/iodata/test/test_prepare.py
+++ b/iodata/test/test_prepare.py
@@ -21,13 +21,16 @@
 from importlib.resources import as_file, files
 from pathlib import Path
 
+import numpy as np
 import pytest
 from numpy.testing import assert_allclose, assert_equal
 
 from ..api import dump_one, load_one
+from ..basis import MolecularBasis, Shell
+from ..convert import HORTON2_CONVENTIONS
 from ..iodata import IOData
 from ..orbitals import MolecularOrbitals
-from ..prepare import prepare_unrestricted_aminusb
+from ..prepare import prepare_segmented, prepare_unrestricted_aminusb
 from ..utils import PrepareDumpError, PrepareDumpWarning
 
 
@@ -104,3 +107,61 @@ def test_dump_occs_aminusb(tmpdir, fmt):
     assert data2.mo.kind == "unrestricted"
     assert_allclose(data2.mo.occsa, data1.mo.occsa)
     assert_allclose(data2.mo.occsb, data1.mo.occsb)
+
+
+def test_segmented_no_basis():
+    data = IOData()
+    with pytest.raises(ValueError):
+        prepare_segmented(data, False, False, "foo.wfn", "wfn")
+
+
+def test_segmented_not_generalized():
+    data = IOData(
+        obasis=MolecularBasis(
+            [
+                Shell(0, [0], ["c"], [0.5, 0.01], [[0.1], [0.2]]),
+                Shell(1, [2], ["p"], [1.1], [[0.3]]),
+            ],
+            HORTON2_CONVENTIONS,
+            "L2",
+        )
+    )
+    assert data is prepare_segmented(data, False, False, "foo.wfn", "wfn")
+
+
+def test_segmented_generalized():
+    rng = np.random.default_rng(1)
+    data0 = IOData(
+        obasis=MolecularBasis(
+            [
+                Shell(0, [0, 1], ["c", "c"], rng.uniform(0, 1, 3), rng.uniform(-1, 1, (3, 2))),
+                Shell(1, [2, 3], ["p", "p"], rng.uniform(0, 1, 4), rng.uniform(-1, 1, (4, 2))),
+            ],
+            HORTON2_CONVENTIONS,
+            "L2",
+        )
+    )
+    with pytest.raises(PrepareDumpError):
+        prepare_segmented(data0, False, False, "foo.wfn", "wfn")
+    with pytest.warns(PrepareDumpWarning):
+        data1 = prepare_segmented(data0, False, True, "foo.wfn", "wfn")
+    assert len(data1.obasis.shells) == 4
+
+
+def test_segmented_sp():
+    rng = np.random.default_rng(1)
+    data0 = IOData(
+        obasis=MolecularBasis(
+            [
+                Shell(0, [0, 1], ["c", "c"], rng.uniform(0, 1, 3), rng.uniform(-1, 1, (3, 2))),
+                Shell(1, [2, 3], ["p", "p"], rng.uniform(0, 1, 4), rng.uniform(-1, 1, (4, 2))),
+            ],
+            HORTON2_CONVENTIONS,
+            "L2",
+        )
+    )
+    with pytest.raises(PrepareDumpError):
+        prepare_segmented(data0, True, False, "foo.wfn", "wfn")
+    with pytest.warns(PrepareDumpWarning):
+        data1 = prepare_segmented(data0, True, True, "foo.wfn", "wfn")
+    assert len(data1.obasis.shells) == 3
diff --git a/iodata/test/test_wfn.py b/iodata/test/test_wfn.py
index 8735625e..26ae3b3f 100644
--- a/iodata/test/test_wfn.py
+++ b/iodata/test/test_wfn.py
@@ -28,8 +28,13 @@
 from ..api import dump_one, load_one
 from ..formats.wfn import load_wfn_low
 from ..overlap import compute_overlap
-from ..utils import LineIterator, PrepareDumpError
-from .common import check_orthonormal, compare_mols, compute_mulliken_charges, create_generalized
+from ..utils import LineIterator, PrepareDumpError, PrepareDumpWarning
+from .common import (
+    check_orthonormal,
+    compare_mols,
+    compute_mulliken_charges,
+    create_generalized_orbitals,
+)
 
 # TODO: removed density, kin, nucnuc checks
 
@@ -300,7 +305,7 @@ def test_load_one_cah110_hf_sto3g_g09():
     check_orthonormal(mol.mo.coeffsb, olp, 1e-5)
 
 
-def check_load_dump_consistency(fn: str, tmpdir: str, atol: float = 1.0e-6):
+def check_load_dump_consistency(fn: str, tmpdir: str, atol: float = 1.0e-6, allow_changes=False):
     """Check if data is preserved after dumping and loading a WFN file.
 
     Parameters
@@ -316,7 +321,7 @@ def check_load_dump_consistency(fn: str, tmpdir: str, atol: float = 1.0e-6):
     with as_file(files("iodata.test.data").joinpath(fn)) as file_name:
         mol1 = load_one(str(file_name))
     fn_tmp = os.path.join(tmpdir, "foo.wfn")
-    dump_one(mol1, fn_tmp)
+    dump_one(mol1, fn_tmp, allow_changes=allow_changes)
     mol2 = load_one(fn_tmp)
     # compare Mulliken charges
     charges1 = compute_mulliken_charges(mol1)
@@ -357,7 +362,8 @@ def test_load_dump_consistency(path, tmpdir):
 
 
 def test_load_dump_consistency_from_fchk_h2o(tmpdir):
-    check_load_dump_consistency("h2o_sto3g.fchk", tmpdir)
+    with pytest.warns(PrepareDumpWarning):
+        check_load_dump_consistency("h2o_sto3g.fchk", tmpdir, allow_changes=True)
 
 
 def test_load_dump_consistency_from_molden_nh3(tmpdir):
@@ -369,8 +375,8 @@ def test_dump_one_pure_functions(tmpdir):
         check_load_dump_consistency("water_ccpvdz_pure_hf_g03.fchk", tmpdir)
 
 
-def test_generalized():
+def test_generalized_orbitals():
     # The WFN format does not support generalized MOs
-    data = create_generalized()
+    data = create_generalized_orbitals()
     with pytest.raises(PrepareDumpError):
         dump_one(data, "generalized.wfn")
diff --git a/iodata/test/test_wfx.py b/iodata/test/test_wfx.py
index 0fb0f394..6f6619e0 100644
--- a/iodata/test/test_wfx.py
+++ b/iodata/test/test_wfx.py
@@ -29,12 +29,12 @@
 from ..api import dump_one, load_one
 from ..formats.wfx import load_data_wfx, parse_wfx
 from ..overlap import compute_overlap
-from ..utils import LineIterator, LoadError, PrepareDumpError
+from ..utils import LineIterator, LoadError, PrepareDumpError, PrepareDumpWarning
 from .common import (
     check_orthonormal,
     compare_mols,
     compute_mulliken_charges,
-    create_generalized,
+    create_generalized_orbitals,
     load_one_warning,
     truncated_file,
 )
@@ -89,7 +89,12 @@ def test_load_dump_consistency_lih_cation_rohf(tmpdir):
 
 
 def compare_mulliken_charges(
-    fname: str, tmpdir: str, rtol: float = 1.0e-7, atol: float = 0.0, match: Optional[str] = None
+    fname: str,
+    tmpdir: str,
+    rtol: float = 1.0e-7,
+    atol: float = 0.0,
+    match: Optional[str] = None,
+    allow_changes: bool = False,
 ):
     """Check if charges are computed correctly after dumping and loading WFX file format.
 
@@ -111,7 +116,11 @@ def compare_mulliken_charges(
     mol1 = load_one_warning(fname, match=match)
     # dump WFX and check that file exists
     fn_tmp = os.path.join(tmpdir, f"{fname}.wfx")
-    dump_one(mol1, fn_tmp)
+    if allow_changes:
+        with pytest.warns(PrepareDumpWarning):
+            dump_one(mol1, fn_tmp, allow_changes=True)
+    else:
+        dump_one(mol1, fn_tmp)
     assert os.path.isfile(fn_tmp)
     # load dumped file and compare Mulliken charges
     mol2 = load_one(fn_tmp)
@@ -156,7 +165,7 @@ def compare_mulliken_charges(
     ],
 )
 def test_dump_one(path, tmpdir):
-    compare_mulliken_charges(path, tmpdir)
+    compare_mulliken_charges(path, tmpdir, allow_changes=path.endswith("fchk"))
 
 
 @pytest.mark.parametrize(
@@ -812,8 +821,8 @@ def test_load_one_lih_cation_rohf():
     check_orthonormal(mol.mo.coeffsb, olp, 1e-5)
 
 
-def test_generalized():
+def test_generalized_orbitals():
     # The Molden format does not support generalized MOs
-    data = create_generalized()
+    data = create_generalized_orbitals()
     with pytest.raises(PrepareDumpError):
         dump_one(data, "generalized.wfx")

From 62dc5ddabda9d3b7c9abb2a29e8e5cee51f6ca7b Mon Sep 17 00:00:00 2001
From: Toon Verstraelen <Toon.Verstraelen@UGent.be>
Date: Thu, 11 Jul 2024 11:52:42 +0200
Subject: [PATCH 2/2] AI-inspired changes

---
 iodata/prepare.py           |  4 ++--
 iodata/test/test_convert.py |  8 ++++++++
 iodata/test/test_molden.py  | 11 +++++------
 3 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/iodata/prepare.py b/iodata/prepare.py
index eafc2d5d..52abcd8c 100644
--- a/iodata/prepare.py
+++ b/iodata/prepare.py
@@ -55,7 +55,7 @@ def prepare_unrestricted_aminusb(data: IOData, allow_changes: bool, filename: st
     -------
     data
         The given data object if no conversion took place,
-        or a shallow copy with some new attriubtes.
+        or a shallow copy with some new attributes.
 
     Raises
     ------
@@ -113,7 +113,7 @@ def prepare_segmented(data: IOData, keep_sp: bool, allow_changes: bool, filename
     -------
     data
         The given data object if no conversion took place,
-        or a shallow copy with some new attriubtes.
+        or a shallow copy with some new attributes.
 
     Raises
     ------
diff --git a/iodata/test/test_convert.py b/iodata/test/test_convert.py
index 2afa2d17..d9b35a86 100644
--- a/iodata/test/test_convert.py
+++ b/iodata/test/test_convert.py
@@ -355,3 +355,11 @@ def test_convert_to_segmented_sp():
     assert_equal(shell3.kinds, ["p"])
     assert_equal(shell3.exponents, obasis0.shells[1].exponents)
     assert_equal(shell3.coeffs, obasis0.shells[1].coeffs[:, 1:])
+
+
+def test_convert_to_segmented_empty():
+    obasis0 = MolecularBasis([], HORTON2_CONVENTIONS, "L2")
+    obasis1 = convert_to_segmented(obasis0, keep_sp=False)
+    assert len(obasis1.shells) == 0
+    obasis2 = convert_to_segmented(obasis0, keep_sp=True)
+    assert len(obasis2.shells) == 0
diff --git a/iodata/test/test_molden.py b/iodata/test/test_molden.py
index 69803710..55a8606e 100644
--- a/iodata/test/test_molden.py
+++ b/iodata/test/test_molden.py
@@ -590,12 +590,11 @@ def test_load_molden_f():
     ],
 )
 def test_load_dump_consistency(tmpdir, fn, match, allow_changes):
-    with as_file(files("iodata.test.data").joinpath(fn)) as file_name:
-        if match is None:
-            mol1 = load_one(str(file_name))
-        else:
-            with pytest.warns(LoadWarning, match=match):
-                mol1 = load_one(str(file_name))
+    with ExitStack() as stack:
+        file_name = stack.enter_context(as_file(files("iodata.test.data").joinpath(fn)))
+        if match is not None:
+            stack.enter_context(pytest.warns(LoadWarning, match=match))
+        mol1 = load_one(file_name)
     fn_tmp = os.path.join(tmpdir, "foo.bar")
     if allow_changes:
         with pytest.warns(PrepareDumpWarning):