OpticSim/mpio.cpp at main · mpoeverlein/OpticSim · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
#include <string>
#include <vector>
#include <sstream>
#include "ray.hpp"
#include "geometry_loader.hpp"
#include "optdev.hpp"

/**
 * @return Python import statements as a string
 */
std::string printImports() {
    std::ostringstream oss;
    oss << "import numpy as np\n"
    << "import matplotlib.pyplot as plt\n"
    << "import matplotlib.patches as patches\n"
    << "import matplotlib as mpl\n";
    return oss.str();
}

/**
 * @return string for Python script to plot rays and optical devices.
 */
std::string printGeometry2D(GeometryLoader& geometry) {
    std::ostringstream oss;
    oss << printImports();
    oss << "mpl.rcParams['lines.linewidth'] = 0.4\n"
    << "fig, ax = plt.subplots()\n";
    for (const auto& device : geometry.devices) {
        oss << device->forPythonPlot();
    }
    for (Ray& r: geometry.rays) {
        oss << r.forPythonPlot();
    }
    oss << "ax.set_aspect('equal', adjustable='box')\n"
    << "plt.show()\n";
    return oss.str();
}

/**
 * @return string for Python script of plotting rays.
 */
std::string printRays(std::vector<Ray> rays) {
    std::ostringstream oss;
    oss << printImports();
    oss << "mpl.rcParams['lines.linewidth'] = 0.4\n"
    << "fig, ax = plt.subplots()\n";

    for (Ray& r: rays) {
        oss << r.forPythonPlot();
    }
    oss << "# ax.set_xlim((8.9,11.1))\n"
        << "# ax.set_ylim((-1.1,1.1))\n"
        << "ax.set_aspect('equal', adjustable='box')\n";
    oss << "plt.show()\n";

    return oss.str();
}