Add program code generator function (#496)

* Add serialize function

* Run black

* Remove f

* Updates from code review

* Improve operations handling

* Remove duplicate line

* Fix generate_code + add tests

* Tidy things up a bit

* Add forgotten factor

* Run black

* Update changelog

* Fixes from code review

* Update strawberryfields/io.py

Co-authored-by: Josh Izaac <josh146@gmail.com>

* change argument name

Co-authored-by: Josh Izaac <josh146@gmail.com>

.github/CHANGELOG.md

@@ -5,6 +5,10 @@
 * `TDMProgram` objects can now be compiled and submitted via the API.
   [(#476)](https://github.com/XanaduAI/strawberryfields/pull/476)
 
+* Strawberry Fields code can be generated from a program (and an engine) by
+  calling `sf.io.generate_code(program, eng=engine)`.
+  [(#496)](https://github.com/XanaduAI/strawberryfields/pull/496)
+
 <h3>Improvements</h3>
 
 * The `copies` option when constructing a `TDMProgram` have been removed. Instead, the number of

strawberryfields/io.py

@@ -12,11 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-This module contains functions for loading and saving Strawberry
-Fields :class:`~.Program` objects from/to Blackbird scripts.
+This module contains functions for loading and saving Strawberry Fields
+:class:`~.Program` objects from/to Blackbird scripts and Strawberry Fields
+code.
 """
 # pylint: disable=protected-access,too-many-nested-blocks
 import os
+from numbers import Number
 
 import numpy as np
 
@@ -244,6 +246,145 @@ def is_free_param(param):
     return prog
 
 
+def generate_code(prog, eng=None):
+    """Converts a Strawberry Fields program into valid Strawberry Fields code.
+
+    **Example:**
+
+    .. code-block:: python3
+
+        prog = sf.Program(3)
+        eng = sf.Engine("fock", backend_options={"cutoff_dim": 5})
+
+        with prog.context as q:
+            ops.Sgate(2*np.pi/3) | q[1]
+            ops.BSgate(0.6, 0.1) | (q[2], q[0])
+            ops.MeasureFock() | q
+
+        results = eng.run(prog)
+
+        code_str = sf.io.generate_code(prog, eng=eng)
+
+    This will create the following string:
+
+    .. code-block:: pycon
+
+        >>> print(code_str)
+        import strawberryfields as sf
+        from strawberryfields import ops
+
+        prog = sf.Program(3)
+        eng = sf.Engine("fock", backend_options={"cutoff_dim": 5})
+
+        with prog.context as q:
+            ops.Sgate(2*np.pi/3, 0.0) | q[1]
+            ops.BSgate(0.6, 0.1) | (q[2], q[0])
+            ops.MeasureFock() | (q[0], q[1], q[2])
+
+        results = eng.run(prog)
+
+    Args:
+        prog (Program): the Strawberry Fields program
+        eng (Engine): The Strawberryfields engine. If ``None``, only the Program
+            parts will be in the resulting code-string.
+
+    Returns:
+        str: the Strawberry Fields code, for constructing the program, as a string
+    """
+    code_seq = ["import strawberryfields as sf", "from strawberryfields import ops\n"]
+
+    if prog.type == "tdm":
+        code_seq.append(f"prog = sf.TDMProgram(N={prog.N})")
+    else:
+        code_seq.append(f"prog = sf.Program({prog.num_subsystems})")
+
+    # check if an engine is supplied; if so, format and add backend/target
+    # along with backend options
+    if eng:
+        eng_type = eng.__class__.__name__
+        if eng_type == "RemoteEngine":
+            code_seq.append(f'eng = sf.RemoteEngine("{eng.target}")')
+        else:
+            if "cutoff_dim" in eng.backend_options:
+                formatting_str = (
+                    f'"{eng.backend_name}", backend_options='
+                    + f'{{"cutoff_dim": {eng.backend_options["cutoff_dim"]}}}'
+                )
+                code_seq.append(f"eng = sf.Engine({formatting_str})")
+            else:
+                code_seq.append(f'eng = sf.Engine("{eng.backend_name}")')
+
+    # check if program is of TDM type and format the context as appropriate
+    if prog.type == "tdm":
+        # if the context arrays contain pi-values, factor out multiples of np.pi
+        tdm_params = [f"[{_factor_out_pi(par)}]" for par in prog.tdm_params]
+        code_seq.append("\nwith prog.context(" + ", ".join(tdm_params) + ") as (p, q):")
+    else:
+        code_seq.append("\nwith prog.context as q:")
+
+    # add the operations, and replace any free parameters with e.g. `p[0]`, `p[1]`
+    for cmd in prog.circuit:
+        name = cmd.op.__class__.__name__
+        if prog.type == "tdm":
+            format_dict = {k: f"p[{k[1:]}]" for k in prog.parameters.keys()}
+            params_str = _factor_out_pi(cmd.op.p).format(**format_dict)
+        else:
+            params_str = _factor_out_pi(cmd.op.p)
+
+        modes = [f"q[{r.ind}]" for r in cmd.reg]
+        if len(modes) == 1:
+            modes_str = ", ".join(modes)
+        else:
+            modes_str = "(" + ", ".join(modes) + ")"
+        op = f"    ops.{name}({params_str}) | {modes_str}"
+
+        code_seq.append(op)
+
+    if eng:
+        code_seq.append("\nresults = eng.run(prog)")
+
+    return "\n".join(code_seq)
+
+
+def _factor_out_pi(num_list, denominator=12):
+    """Factors out pi, divided by the denominator value, from all number in a list
+    and returns a string representation.
+
+    Args:
+        num_list (list[Number, string]): a list of numbers and/or strings
+        denominator (int): factor out pi divided by denominator;
+            e.g. default would be to factor out np.pi/12
+
+    Return:
+        string: containing strings of values and/or input string objects
+    """
+    factor = np.pi / denominator
+
+    a = []
+    for p in num_list:
+        # if list element is not a number then append its string representation
+        if not isinstance(p, Number):
+            a.append(str(p))
+            continue
+
+        if np.isclose(p % factor, [0, factor]).any() and p != 0:
+            gcd = np.gcd(int(p / factor), denominator)
+            if gcd == denominator:
+                if int(p / np.pi) == 1:
+                    a.append("np.pi")
+                else:
+                    a.append(f"{int(p/np.pi)}*np.pi")
+            else:
+                coeff = int(p / factor / gcd)
+                if coeff == 1:
+                    a.append(f"np.pi/{int(denominator/gcd)}")
+                else:
+                    a.append(f"{coeff}*np.pi/{int(denominator / gcd)}")
+        else:
+            a.append(str(p))
+    return ", ".join(a)
+
+
 def save(f, prog):
     """Saves a quantum program to a Blackbird .xbb file.
 

tests/frontend/test_io.py

@@ -644,6 +644,155 @@ def test_gate_not_defined_tdm(self):
         with pytest.raises(NameError, match="operation np not defined"):
             io.to_program(bb)
 
+
+prog_txt = textwrap.dedent('''\
+    import strawberryfields as sf
+    from strawberryfields import ops
+
+    prog = sf.Program(3)
+    eng = sf.Engine({engine_args})
+
+    with prog.context as q:
+        ops.Sgate(0.54, 0) | q[0]
+        ops.BSgate(0.45, np.pi/2) | (q[0], q[2])
+        ops.Sgate(3*np.pi/2, 0) | q[1]
+        ops.BSgate(2*np.pi, 0.62) | (q[0], q[1])
+        ops.MeasureFock() | q[0]
+
+    results = eng.run(prog)
+''')
+
+prog_txt_no_engine = textwrap.dedent('''\
+    import strawberryfields as sf
+    from strawberryfields import ops
+
+    prog = sf.Program(3)
+
+    with prog.context as q:
+        ops.Sgate(0.54, 0) | q[0]
+        ops.BSgate(0.45, np.pi/2) | (q[0], q[2])
+        ops.Sgate(3*np.pi/2, 0) | q[1]
+        ops.BSgate(2*np.pi, 0.62) | (q[0], q[1])
+        ops.MeasureFock() | q[0]
+''')
+
+prog_txt_tdm = textwrap.dedent('''\
+    import strawberryfields as sf
+    from strawberryfields import ops
+
+    prog = sf.TDMProgram(N=[2, 3])
+    eng = sf.Engine("gaussian")
+
+    with prog.context([np.pi, 3*np.pi/2, 0], [1, 0.5, np.pi], [0, 0, 0]) as (p, q):
+        ops.Sgate(0.123, np.pi/4) | q[2]
+        ops.BSgate(p[0], 0.0) | (q[1], q[2])
+        ops.Rgate(p[1]) | q[2]
+        ops.MeasureHomodyne(p[0]) | q[0]
+        ops.MeasureHomodyne(p[2]) | q[2]
+
+    results = eng.run(prog)
+''')
+
+class TestGenerateCode:
+    """Tests for the generate_code function"""
+
+    def test_generate_code_no_engine(self):
+        """Test generating code for a regular program with no engine"""
+        prog = sf.Program(3)
+
+        with prog.context as q:
+            ops.Sgate(0.54, 0) | q[0]
+            ops.BSgate(0.45, np.pi/2) | (q[0], q[2])
+            ops.Sgate(3*np.pi/2, 0) | q[1]
+            ops.BSgate(2*np.pi, 0.62) | (q[0], q[1])
+            ops.MeasureFock() | q[0]
+
+        code = io.generate_code(prog)
+
+        code_list = code.split("\n")
+        expected = prog_txt_no_engine.split("\n")
+
+        for i, row in enumerate(code_list):
+            assert row == expected[i]
+
+    @pytest.mark.parametrize("engine_kwargs", [
+        {"backend": "fock", "backend_options": {"cutoff_dim": 5}},
+        {"backend": "gaussian"},
+    ])
+    def test_generate_code_with_engine(self, engine_kwargs):
+        """Test generating code for a regular program with an engine"""
+        prog = sf.Program(3)
+        eng = sf.Engine(**engine_kwargs)
+
+        with prog.context as q:
+            ops.Sgate(0.54, 0) | q[0]
+            ops.BSgate(0.45, np.pi/2) | (q[0], q[2])
+            ops.Sgate(3*np.pi/2, 0) | q[1]
+            ops.BSgate(2*np.pi, 0.62) | (q[0], q[1])
+            ops.MeasureFock() | q[0]
+
+        results = eng.run(prog)
+
+        code = io.generate_code(prog, eng)
+
+        code_list = code.split("\n")
+        formatting_str = f"\"{engine_kwargs['backend']}\""
+        if "backend_options" in engine_kwargs:
+            formatting_str += (
+                ", backend_options="
+                f'{{"cutoff_dim": {engine_kwargs["backend_options"]["cutoff_dim"]}}}'
+            )
+        expected = prog_txt.format(engine_args=formatting_str).split("\n")
+
+        for i, row in enumerate(code_list):
+            assert row == expected[i]
+
+
+    def test_generate_code_tdm(self):
+        """Test generating code for a TDM program with an engine"""
+        prog = sf.TDMProgram(N=[2, 3])
+        eng = sf.Engine("gaussian")
+
+        with prog.context([np.pi, 3*np.pi/2, 0], [1, 0.5, np.pi], [0, 0, 0]) as (p, q):
+            ops.Sgate(0.123, np.pi/4) | q[2]
+            ops.BSgate(p[0]) | (q[1], q[2])
+            ops.Rgate(p[1]) | q[2]
+            ops.MeasureHomodyne(p[0]) | q[0]
+            ops.MeasureHomodyne(p[2]) | q[2]
+
+        results = eng.run(prog)
+
+        code = io.generate_code(prog, eng)
+
+        code_list = code.split("\n")
+        expected = prog_txt_tdm.split("\n")
+
+        for i, row in enumerate(code_list):
+            assert row == expected[i]
+
+    @pytest.mark.parametrize(
+        "value", [
+            "np.pi", "np.pi/2", "np.pi/12", "2*np.pi", "2*np.pi/3",
+            "0", "42", "9.87", "-0.2", "no_number", "{p3}"
+        ],
+    )
+    def test_factor_out_pi(self, value):
+        """Test that the factor_out_pi function is able to convert floats
+        that are equal to a pi expression, to strings containing a pi
+        expression.
+
+        For example, the float 6.28318530718 should be converted to the string "2*np.pi"
+        """
+        try:
+            val = eval(value)
+        except NameError:
+            val = value
+        res = sf.io._factor_out_pi([val])
+        expected = value
+
+        assert res == expected
+
+
 class DummyResults:
     """Dummy results object"""