Task: fix supercell with symmetry (pyscf/pyscf#1773)

org: pyscf
repo: pyscf
pr_number: 1773
issue_title: fix supercell with symmetry
issue_body: Fix issue #1758.
resolved_issues: number: 1758
title: Error when calculating the IAOs of a supercell with space_group_symmetry=True.
body: Hello!
I've come across an error when trying to obtain the IAOs of a supercell generated from a cell initialised with space group symmetry.
A minimal example reproducing the issue is
```
import pyscf.lo
import pyscf.pbc
import numpy as np
​
cell = pyscf.pbc.gto.Cell()
cell.atom = ['He 0.0 0.0 0.0']
cell.a = 1.4 * np.eye(3)
cell.basis = 'def2-svp'
cell.output = 'pyscf.out'
cell.space_group_symmetry = True
cell.symmorphic = True
cell.build()
​
scell = pyscf.pbc.tools.super_cell(cell, (2,2,2))
pyscf.lo.iao.reference_mol(scell, minao="sto-3g")
```

which gives traceback

```
Traceback (most recent call last):
 File "testscript.py", line 16, in <module>
    pyscf.lo.iao.reference_mol(scell, minao="sto-3g")
 File "python3.10/site-packages/pyscf/lo/iao.py", line 135, in reference_mol
    pmol.build(False, False, basis=minao)
 File "python3.10/site-packages/pyscf/pbc/gto/cell.py", line 1516, in build
    _lattice_symm.build(space_group_symmetry=True,
 File "python3.10/site-packages/pyscf/pbc/symm/symmetry.py", line 176, in build
    self.spacegroup = space_group.SpaceGroup(cell).build(dump_info=False)
 File "python3.10/site-packages/pyscf/pbc/symm/space_group.py", line 302, in build
    self.ops = geom.search_space_group_ops(self.cell, tol=self.symprec)
 File "python3.10/site-packages/pyscf/pbc/symm/geom.py", line 81, in search_space_group_ops
    atmgrp = mole.atom_types(cell._atom, magmom=cell.magmom)
 File "python3.10/site-packages/pyscf/gto/mole.py", line 308, in atom_types
    uniq_mag = np.unique(magmom[idx])
IndexError: index 1 is out of bounds for axis 0 with size 1
```

This can be avoided by disabling space group symmetry initially or in the supercell object after generation, but I don't know whether that should be necessary. Any advice as to whether this is an intended incompatibility would be greatly appreciated.


Reference solution:
fix_patch: diff --git a/pyscf/gto/mole.py b/pyscf/gto/mole.py
index 97e3c650ba..1d832b2b43 100644
--- a/pyscf/gto/mole.py
+++ b/pyscf/gto/mole.py
@@ -304,7 +304,7 @@ def atom_types(atoms, basis=None, magmom=None):
 
     if magmom is not None:
         atmgroup_new = {}
-        suffix = {1.0:'u', -1.0:'d', 0.0:'o'}
+        suffix = {1:'u', -1:'d', 0:'o'}
         magmom = np.asarray(magmom)
         for elem, idx in atmgroup.items():
             uniq_mag = np.unique(magmom[idx])
@@ -312,7 +312,7 @@ def atom_types(atoms, basis=None, magmom=None):
                 for i, mag in enumerate(uniq_mag):
                     subgrp = np.asarray(idx)[np.where(magmom[idx] == mag)[0]]
                     if mag not in suffix:
-                        raise RuntimeError("Magmom should be chosen from [-1., 0., 1.], but %s is given" % mag)
+                        raise RuntimeError("Magmom should be chosen from [-1, 0, 1], but %s is given" % mag)
                     atmgroup_new[elem+'_'+suffix[mag]] = subgrp.tolist()
             else:
                 atmgroup_new[elem] = idx
@@ -2160,7 +2160,7 @@ class Mole(lib.StreamObject):
         cart : boolean
             Using Cartesian GTO basis and integrals (6d,10f,15g)
         magmom : list
-            Collinear spin of each atom. Default is [0.0,]*natm
+            Collinear spin of each atom. Default is [0,]*natm
 
         ** Following attributes are generated by :func:`Mole.build` **
 
@@ -2259,7 +2259,7 @@ def __init__(self, **kwargs):
         self.ecp = {}
 # Nuclear property. self.nucprop = {atom_symbol: {key: value}}
         self.nucprop = {}
-# Collinear spin of each atom. self.magmom = [0., ...]
+# Collinear spin of each atom. self.magmom = [0, ...]
         self.magmom = []
 ##################################################
 # don't modify the following private variables, they are not input options
@@ -2470,7 +2470,7 @@ def build(self, dump_input=True, parse_arg=ARGPARSE,
                 Whether to use symmetry.  If given a string of point group
                 name, the given point group symmetry will be used.
             magmom : list
-                Collinear spin of each atom. Default is [0.0,]*natm
+                Collinear spin of each atom. Default is [0,]*natm
 
         '''
         if not DISABLE_GC:
@@ -2563,9 +2563,9 @@ def build(self, dump_input=True, parse_arg=ARGPARSE,
             # number of electrons are consistent.
             self.nelec
 
-        if not self.magmom:
-            self.magmom = [0.,]*self.natm
-        if self.spin == 0 and abs(numpy.sum(numpy.asarray(self.magmom)) - self.spin) > 1e-6:
+        if self.magmom is None or len(self.magmom) != self.natm:
+            self.magmom = [0,] * self.natm
+        if self.spin == 0 and abs(numpy.sum(self.magmom) - self.spin) > 1e-6:
             #don't check for unrestricted calcs.
             raise ValueError("mol.magmom is set incorrectly.")
 
diff --git a/pyscf/pbc/gto/cell.py b/pyscf/pbc/gto/cell.py
index 075a69e8eb..8ae280772a 100644
--- a/pyscf/pbc/gto/cell.py
+++ b/pyscf/pbc/gto/cell.py
@@ -276,12 +276,7 @@ def byteify(inp):
     # Symmetry class cannot be serialized by dumps function.
     # Recreate it manually
     if cell.natm > 0 and cell.space_group_symmetry:
-        from pyscf.pbc.symm import Symmetry
-        _lattice_symm = Symmetry(cell)
-        _lattice_symm.build(space_group_symmetry=True,
-                            symmorphic=cell.symmorphic,
-                            check_mesh_symmetry=False)
-        cell.lattice_symmetry = _lattice_symm
+        cell.build_lattice_symmetry()
 
     return cell
 
@@ -1268,6 +1263,32 @@ def symmetrize_mesh(self, mesh=None):
             logger.debug(self, 'mesh %s is symmetrized as %s', mesh, mesh1)
         return mesh1
 
+    def build_lattice_symmetry(self, check_mesh_symmetry=True):
+        '''Build cell.lattice_symmetry object.
+
+        Kwargs:
+            check_mesh_symmetry : bool
+                For nonsymmorphic symmetry groups, `cell.mesh` may have
+                lower symmetry than the lattice. In this case, if
+                `check_mesh_symmetry` is `True`, the lower symmetry group will
+                be used. Otherwise, if `check_mesh_symmetry` is `False`,
+                the mesh grid will be modified to satisfy the higher symmetry.
+                Default value is `True`.
+
+        Note:
+            This function modifies the attributes of `cell`.
+        '''
+        from pyscf.pbc.symm import Symmetry
+        self.lattice_symmetry = Symmetry(self).build(
+                                    space_group_symmetry=True,
+                                    symmorphic=self.symmorphic,
+                                    check_mesh_symmetry=check_mesh_symmetry)
+        if not check_mesh_symmetry:
+            _mesh_from_build = self._mesh_from_build
+            self.mesh = self.symmetrize_mesh()
+            self._mesh_from_build = _mesh_from_build
+        return self
+
 #Note: Exculde dump_input, parse_arg, basis from kwargs to avoid parsing twice
     def build(self, dump_input=True, parse_arg=mole.ARGPARSE,
               a=None, mesh=None, ke_cutoff=None, precision=None, nimgs=None,
@@ -1484,17 +1505,8 @@ def build(self, dump_input=True, parse_arg=mole.ARGPARSE,
             self._mesh_from_build = True
 
         if self.space_group_symmetry:
-            from pyscf.pbc.symm import Symmetry
-            _check_mesh_symm = False
-            if not self._mesh_from_build:
-                _check_mesh_symm = True
-            _lattice_symm = Symmetry(self)
-            _lattice_symm.build(space_group_symmetry=True,
-                                symmorphic=self.symmorphic,
-                                check_mesh_symmetry=_check_mesh_symm)
-            self.lattice_symmetry = _lattice_symm
-            if self._mesh_from_build:
-                self._mesh = self.symmetrize_mesh()
+            _check_mesh_symm = not self._mesh_from_build
+            self.build_lattice_symmetry(check_mesh_symmetry=_check_mesh_symm)
 
         if dump_input and not _built and self.verbose > logger.NOTE:
             self.dump_input()
diff --git a/pyscf/pbc/symm/geom.py b/pyscf/pbc/symm/geom.py
index 32b40c44bd..74119a4483 100644
--- a/pyscf/pbc/symm/geom.py
+++ b/pyscf/pbc/symm/geom.py
@@ -224,7 +224,7 @@ def get_crystal_class(cell, ops=None, tol=SYMPREC):
     """
     cell.a = [[4.0, 0., 0.], [0., 4.0, 0.], [0., 0., 16.0]]
     cell.dimension = 2
-    cell.magmom = [1., 1., -1., -1., 1., -1., 1., 1., -1., -1., 1., -1.]
+    cell.magmom = [1, 1, -1, -1, 1, -1, 1, 1, -1, -1, 1, -1]
     cell.build()
 
     ops = search_space_group_ops(cell)
diff --git a/pyscf/pbc/tools/pbc.py b/pyscf/pbc/tools/pbc.py
index 9945c59f7f..985c0f2550 100644
--- a/pyscf/pbc/tools/pbc.py
+++ b/pyscf/pbc/tools/pbc.py
@@ -589,6 +589,8 @@ def super_cell(cell, ncopy, wrap_around=False):
     supcell.a = np.einsum('i,ij->ij', ncopy, a)
     mesh = np.asarray(ncopy) * np.asarray(cell.mesh)
     supcell.mesh = (mesh // 2) * 2 + 1
+    supcell.magmom = np.repeat(np.asarray(cell.magmom).reshape(1,-1),
+                               np.prod(ncopy), axis=0).ravel()
     return _build_supcell_(supcell, cell, Ls)
 
 
@@ -621,6 +623,7 @@ def _build_supcell_(supcell, cell, Ls):
     Construct supcell ._env directly without calling supcell.build() method.
     This reserves the basis contraction coefficients defined in cell
     '''
+    from pyscf.pbc import gto as pbcgto
     nimgs = len(Ls)
     symbs = [atom[0] for atom in cell._atom] * nimgs
     coords = Ls.reshape(-1,1,3) + cell.atom_coords()
@@ -646,6 +649,9 @@ def _build_supcell_(supcell, cell, Ls):
     supcell._atm = np.asarray(_atm, dtype=np.int32)
     supcell._bas = np.asarray(_bas.reshape(-1, BAS_SLOTS), dtype=np.int32)
     supcell._env = _env
+
+    if isinstance(supcell, pbcgto.Cell) and supcell.space_group_symmetry:
+        supcell.build_lattice_symmetry(not cell._mesh_from_build)
     return supcell