amended requests from pfebrer and tfrederiksen

sisl.geom geometries intrinsically sets the PBC

single string for short setting them is now doable

Signed-off-by: Nick Papior <[email protected]>

src/sisl/geom/_common.py

@@ -10,16 +10,22 @@ def geometry_define_nsc(geometry, periodic=(True, True, True)):
     """Define the number of supercells for a geometry based on the periodicity """
     if np.all(geometry.maxR(True) > 0.):
         geometry.optimize_nsc()
-        if not periodic[0]:
-            geometry.set_nsc(a=1)
-        if not periodic[1]:
-            geometry.set_nsc(b=1)
-        if not periodic[2]:
-            geometry.set_nsc(c=1)
+        for i, d, per in zip(range(3), "abc", periodic):
+            if per:
+                geometry.lattice.set_boundary_condition(**{d: "Periodic"})
+            else:
+                geometry.set_nsc(**{d: 1})
     else:
         nsc = [3 if p else 1 for p in periodic]
         geometry.set_nsc(nsc)
 
+    for name, per in zip("abc", periodic):
+        if per:
+            per = "Periodic"
+        else:
+            per = "Unknown"
+        geometry.lattice.set_boundary_condition(**{name: per})
+
 
 def geometry2uc(geometry, dx=1e-8):
     """ Translate the geometry to the unit cell by first shifting `dx` """

src/sisl/lattice.py

@@ -47,12 +47,33 @@ class BoundaryCondition(IntEnum):
     @classmethod
     def getitem(cls, key):
         """Search for a specific integer entry by value, and not by name """
-        for bc in cls:
-            if bc == key:
-                return bc
+        if isinstance(key, cls):
+            return key
+        if isinstance(key, bool):
+            if key:
+                return cls.PERIODIC
+            raise ValueError(f"{cls.__name__}.getitem does not allow False, which BC should this refer to?")
+        if isinstance(key, str):
+            key = key.upper()
+            if len(key) == 1:
+                key = {"U": "UNKNOWN",
+                       "P": "PERIODIC",
+                       "D": "DIRICHLET",
+                       "N": "NEUMANN",
+                       "O": "OPEN",
+                }[key]
+            for bc in cls:
+                if bc.name.startswith(key):
+                    return bc
+        else:
+            for bc in cls:
+                if bc == key:
+                    return bc
         raise KeyError(f"{cls.__name__}.getitem could not find key={key}")
 
-BoundaryConditionType = Union[int, Sequence[int]]
+BoundaryConditionType = Union[BoundaryCondition, int, str, bool]
+SeqBoundaryConditionType = Union[BoundaryConditionType,
+                                 Sequence[BoundaryConditionType]]
 
 
 @set_module("sisl")
@@ -95,7 +116,7 @@ class Lattice(_Dispatchs,
     BC = BoundaryCondition
 
     def __init__(self, cell, nsc=None, origin=None,
-                 boundary_condition: Sequence[BoundaryConditionType] =BoundaryCondition.PERIODIC):
+                 boundary_condition: SeqBoundaryConditionType =BoundaryCondition.PERIODIC):
 
         if nsc is None:
             nsc = [1, 1, 1]
@@ -168,10 +189,10 @@ def toCuboid(self, *args, **kwargs):
         return self.to[Cuboid](*args, **kwargs)
 
     def set_boundary_condition(self,
-                               boundary: Optional[Sequence[BoundaryConditionType]] =None,
-                               a: Sequence[BoundaryConditionType] =None,
-                               b: Sequence[BoundaryConditionType] =None,
-                               c: Sequence[BoundaryConditionType] =None):
+                               boundary: Optional[SeqBoundaryConditionType] =None,
+                               a: Optional[SeqBoundaryConditionType] =None,
+                               b: Optional[SeqBoundaryConditionType] =None,
+                               c: Opitonal[SeqBoundaryConditionType] =None):
         """ Set the boundary conditions on the grid
 
         Parameters
@@ -190,19 +211,38 @@ def set_boundary_condition(self,
         ValueError
             if specifying periodic one one boundary, so must the opposite side.
         """
+        getitem = BoundaryCondition.getitem
+        def conv(v):
+            if v is None:
+                return v
+            if isinstance(v, (np.ndarray, list, tuple)):
+                return list(map(getitem, v))
+            return getitem(v)
+
+        if not hasattr(self, "_bc"):
+            self._bc = _a.fulli([3, 2], getitem("Unknown"))
+
         if not boundary is None:
-            self._bc = _a.emptyi([3, 2])
-            if isinstance(boundary, Integral):
-                self._bc[:, :] = boundary
+            if isinstance(boundary, (Integral, str, bool)):
+                try:
+                    getitem(boundary)
+                    self._bc[:, :] = conv(boundary)
+                except KeyError:
+                    for d, bc in enumerate(boundary):
+                        bc = conv(bc)
+                        if bc is not None:
+                            self._bc[d] = conv(bc)
+
             else:
-                for i, bc in enumerate(boundary):
-                    self._bc[i] = bc
-        if not a is None:
-            self._bc[0, :] = a
-        if not b is None:
-            self._bc[1, :] = b
-        if not c is None:
-            self._bc[2, :] = c
+                for d, bc in enumerate(boundary):
+                    bc = conv(bc)
+                    if bc is not None:
+                        self._bc[d] = bc
+
+        for d, v in enumerate([a, b, c]):
+            v = conv(v)
+            if v is not None:
+                self._bc[d, :] = v
 
         # shorthand for bc
         for bc in self._bc == BoundaryCondition.PERIODIC:

src/sisl/tests/test_lattice.py

@@ -508,3 +508,34 @@ def test_lattice_bc_set():
     assert not lat.pbc.any()
     assert (lat.boundary_condition == Lattice.BC.UNKNOWN).all()
 
+    lat.boundary_condition = ['per', "unkn", [3, 4]]
+
+    for n in "abc":
+        lat.set_boundary_condition(**{n: "per"})
+        lat.set_boundary_condition(**{n: [3, Lattice.BC.UNKNOWN]})
+        lat.set_boundary_condition(**{n: [True, Lattice.BC.PERIODIC]})
+
+    bc = [
+            "per",
+            ["Dirichlet", Lattice.BC.NEUMANN],
+            ["un", "neu"]
+    ]
+    lat.set_boundary_condition(bc)
+    assert np.all(lat.boundary_condition[1] == [Lattice.BC.DIRICHLET, Lattice.BC.NEUMANN])
+    assert np.all(lat.boundary_condition[2] == [Lattice.BC.UNKNOWN, Lattice.BC.NEUMANN])
+    lat.set_boundary_condition(["per", None, ["dirichlet", "unkno"]])
+    assert np.all(lat.boundary_condition[1] == [Lattice.BC.DIRICHLET, Lattice.BC.NEUMANN])
+    assert np.all(lat.boundary_condition[2] == [Lattice.BC.DIRICHLET, Lattice.BC.UNKNOWN])
+    lat.set_boundary_condition("ppu")
+    assert np.all(lat.pbc == [True, True, False])
+
+
+def test_lattice_bc_fail():
+    lat = Lattice(1)
+    with pytest.raises(ValueError):
+        lat.boundary_condition = False
+    with pytest.raises(ValueError):
+        lat.set_boundary_condition(b=False)
+    with pytest.raises(KeyError):
+        lat.set_boundary_condition(b="eusoatuhesoau")
+