consolidate different pNBDAC to a single class

surface_coatings/chains/nbdac_polymer.py

@@ -81,7 +81,7 @@ def __init__(
         self.labels["down"] = backbone["down"]
 
 
-class SilaneNBDAC(mb.Compound):
+class pNBDAC(mb.Compound):
     """A general method to create a NBDAC polymer with varying side chains/terminal groups.
 
     Parameters
@@ -94,8 +94,11 @@ class SilaneNBDAC(mb.Compound):
         Cap the front of the polymer (NBDAC end)
     cap_end : bool, optional, default=False
         Cap the end of the polymer (Silane end)
-    silane_buffer : int, optional, default=1
-        Silane monomer used to buffer at one end of the NBDAC polymer
+    buffer : str, optional, default=None
+        Type of buffer to be attached to the end of the pNBDAC.
+        Available options: None, "alkyl", "ether", "silane"
+    buffer_length : int, optional, default=1
+        Length of CH2 monomer to be added to the buffer at one end of the NBDAC polymer
     n : int, optional, default=1
         Number of repeat for the monomer
     port_labels: tuple, optional, default=('up', 'down')
@@ -111,14 +114,15 @@ def __init__(
         monomer,
         side_chains=AminoPropyl(),
         terminal_groups=Acetaldehyde(),
-        silane_buffer=1,
+        buffer=None,
+        buffer_length=1,
         cap_front=True,
         cap_end=False,
         n=1,
         port_labels=("up", "down"),
         align=True,
     ):
-        super(SilaneNBDAC, self).__init__()
+        super(pNBDAC, self).__init__()
         if monomer:
             assert isinstance(monomer, fmNBDAC)
             if side_chains and terminal_groups:
@@ -134,12 +138,16 @@ def __init__(
         polymer.build(n=n, add_hydrogens=False)
         self.add(polymer, "Polymer")
 
-        if silane_buffer:
+        if buffer:
             tail = mb.Compound(name="tail")
-            tail.add(Silane(), "Silane")
+            buffer_options = {"alkyl": CH2,
+                              "ether": O,
+                              "silane": Silane}
+            tail.add(buffer_options.get(buffer.lower())(), "Buffer")
             tail.add(CH2(), "CH2_0")
 
-            for i in range(silane_buffer):
+            i = 0
+            for i in range(buffer_length):
                 tail.add(CH2(), f"CH2_{i+1}")
                 mb.force_overlap(
                     tail[f"CH2_{i+1}"],
@@ -148,7 +156,7 @@ def __init__(
                 )
 
             mb.force_overlap(
-                tail["Silane"], tail["Silane"]["up"], tail[f"CH2_0"]["down"]
+                tail["Buffer"], tail["Buffer"]["up"], tail[f"CH2_0"]["down"]
             )
 
             self.add(tail, "tail")
@@ -160,129 +168,6 @@ def __init__(
                 )
                 norm_vector = polymer_vector / np.linalg.norm(polymer_vector)
 
-                new_port = mb.Port(
-                    anchor=polymer[port_labels[0]].anchor,
-                    orientation=norm_vector,
-                    separation=0.07,
-                )
-
-                polymer[port_labels[0]].update_orientation(
-                    orientation=norm_vector
-                )
-
-            mb.force_overlap(
-                tail, tail[f"CH2_{i+1}"]["up"], polymer[port_labels[0]]
-            )
-
-            self.labels["up"] = self["Polymer"][port_labels[1]]
-            self.labels["down"] = self["tail"]["Silane"]["down"]
-        else:
-            self.labels["up"] = self["Polymer"]["up"]
-            self.labels["down"] = self["Polymer"]["down"]
-
-        if cap_front:
-            front_cap = H()
-            self.add(front_cap, "front_cap")
-            mb.force_overlap(
-                move_this=front_cap,
-                from_positions=front_cap["up"],
-                to_positions=self["up"],
-            )
-        if cap_end:
-            end_cap = H()
-            self.add(end_cap, "end_cap")
-            mb.force_overlap(
-                move_this=end_cap,
-                from_positions=end_cap["up"],
-                to_positions=self["down"],
-            )
-
-class EtherNBDAC(mb.Compound):
-    """A general method to create a NBDAC polymer with varying side chains/terminal groups.
-    Option to add an ether cap (to connect to a surface)
-
-    Parameters
-    ----------
-    side_chains : mb.Compound or list of Compounds (len 2)
-        Side chains attached to NBDAC monomer
-    terminal_groups: mb.Compound or list of Compounds (len 2)
-        Terminal groups which will be matched with side chains
-    cap_front : bool, optional, default=True
-        Cap the front of the polymer (NBDAC end)
-    cap_end : bool, optional, default=False
-        Cap the end of the polymer (Silane end)
-    ether_buffer : int, optional, default=1
-        Ether used to buffer at one end of the NBDAC polymer
-    n : int, optional, default=1
-        Number of repeat for the monomer
-    port_labels: tuple, optional, default=('up', 'down')
-        The list of ports of the monomers. The Silane will connect with the first port
-        and the last port will be capped by a Hydrogen
-    algin : bool, optional, default=True
-        If True, align the port connected to the silane buffer with the rest of the polymer.
-        The goal is to create a straight polymer when grafted on a surface.
-    """
-
-    def __init__(
-        self,
-        monomer,
-        side_chains=AminoPropyl(),
-        terminal_groups=Acetaldehyde(),
-        ether_buffer=1,
-        cap_front=True,
-        cap_end=False,
-        n=1,
-        port_labels=("up", "down"),
-        align=True,
-    ):
-        super(EtherNBDAC, self).__init__()
-        if monomer:
-            assert isinstance(monomer, fmNBDAC)
-            if side_chains and terminal_groups:
-                warnings.warn(
-                    "Both monomer and (side_chains, terminal_groups) are provided,"
-                    "only monomer will be used to construct the polymer."
-                    "Please refer to docstring for more information."
-                )
-        else:
-            monomer = fmNBDAC(side_chains, terminal_groups)
-
-        polymer = Polymer(monomers=[monomer])
-        polymer.build(n=n, add_hydrogens=False)
-        self.add(polymer, "Polymer")
-
-        if ether_buffer:
-            tail = mb.Compound(name="tail")
-            tail.add(O(), "O_init")
-            tail.add(CH2(), "CH2_0")
-
-            for i in range(ether_buffer):
-                tail.add(CH2(), f"CH2_{i+1}")
-                mb.force_overlap(
-                    tail[f"CH2_{i+1}"],
-                    tail[f"CH2_{i+1}"]["down"],
-                    tail[f"CH2_{i}"]["up"],
-                )
-
-            mb.force_overlap(
-                tail["O_init"], tail["O_init"]["up"], tail[f"CH2_0"]["down"]
-            )
-
-            self.add(tail, "tail")
-
-            if align:
-                polymer_vector = (
-                    polymer[port_labels[1]].anchor.pos
-                    - polymer[port_labels[0]].anchor.pos
-                )
-                norm_vector = polymer_vector / np.linalg.norm(polymer_vector)
-
-                new_port = mb.Port(
-                    anchor=polymer[port_labels[0]].anchor,
-                    orientation=norm_vector,
-                    separation=0.07,
-                )
-
                 polymer[port_labels[0]].update_orientation(
                     orientation=norm_vector
                 )
@@ -292,7 +177,7 @@ def __init__(
             )
 
             self.labels["up"] = self["Polymer"][port_labels[1]]
-            self.labels["down"] = self["tail"]["O_init"]["down"]
+            self.labels["down"] = self["tail"]["Buffer"]["down"]
         else:
             self.labels["up"] = self["Polymer"]["up"]
             self.labels["down"] = self["Polymer"]["down"]