290 lines
9.3 KiB
Python
290 lines
9.3 KiB
Python
|
"""
|
||
|
|
Extension for rendering beams in 2D optics with Inkscape
|
||
|
|
"""
|
||
|
|
from __future__ import annotations
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from functools import singledispatchmethod
|
||
|
|
from typing import Iterable, Optional, Final
|
||
|
|
import inkex
|
||
|
|
from inkex.paths import Line, Move
|
||
|
|
import raytracing.material
|
||
|
|
from desc_parser import get_optics_fields
|
||
|
|
from raytracing import Vector, World, OpticalObject, Ray
|
||
|
|
from raytracing.geometry import CubicBezier, CompoundGeometricObject, GeometricObject
|
||
|
|
from utils import pairwise
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class BeamSeed:
|
||
|
|
ray: Optional[Ray] = None
|
||
|
|
parent: Optional[inkex.ShapeElement] = None
|
||
|
|
|
||
|
|
|
||
|
|
def get_unlinked_copy(clone: inkex.Use) -> Optional[inkex.ShapeElement]:
|
||
|
|
"""Creates a copy of the original with all transformations applied"""
|
||
|
|
copy = clone.href.copy()
|
||
|
|
copy.transform = clone.composed_transform() * copy.transform
|
||
|
|
copy.style = clone.style + copy.style
|
||
|
|
copy.getparent = clone.getparent
|
||
|
|
return copy
|
||
|
|
|
||
|
|
|
||
|
|
def get_or_create_beam_layer(parent_layer: inkex.Layer) -> inkex.Layer:
|
||
|
|
for child in parent_layer:
|
||
|
|
if isinstance(child, inkex.Layer):
|
||
|
|
if child.get("inkscape:label") == "generated_beams":
|
||
|
|
return child
|
||
|
|
new_layer = parent_layer.add(inkex.Layer())
|
||
|
|
new_layer.label = "generated_beams"
|
||
|
|
return new_layer
|
||
|
|
|
||
|
|
|
||
|
|
def plot_beam(beam: list[Ray], node: inkex.ShapeElement, layer: inkex.Layer):
|
||
|
|
path = inkex.Path()
|
||
|
|
if beam:
|
||
|
|
path += [Move(beam[0].origin.x, beam[0].origin.y)]
|
||
|
|
for ray in beam:
|
||
|
|
p1 = ray.origin + ray.travel * ray.direction
|
||
|
|
path += [Line(p1.x, p1.y)]
|
||
|
|
element = layer.add(inkex.PathElement())
|
||
|
|
# Need to convert to path to get the correct style for inkex.Use
|
||
|
|
element.style = node.to_path_element().style
|
||
|
|
element.path = path
|
||
|
|
|
||
|
|
|
||
|
|
class Raytracing(inkex.EffectExtension):
|
||
|
|
"""Extension to renders the beams present in the document"""
|
||
|
|
|
||
|
|
# Ray tracing is only implemented for the following inkex primitives
|
||
|
|
filter_primitives: Final = (
|
||
|
|
inkex.PathElement,
|
||
|
|
inkex.Line,
|
||
|
|
inkex.Polyline,
|
||
|
|
inkex.Polygon,
|
||
|
|
inkex.Rectangle,
|
||
|
|
inkex.Ellipse,
|
||
|
|
inkex.Circle,
|
||
|
|
)
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
super().__init__()
|
||
|
|
self.world = World()
|
||
|
|
self.beam_seeds: list[BeamSeed] = list()
|
||
|
|
|
||
|
|
def effect(self) -> None:
|
||
|
|
"""
|
||
|
|
Loads the objects and outputs a svg with the beams after propagation
|
||
|
|
"""
|
||
|
|
|
||
|
|
# Can't set the border earlier because self.svg is not yet defined
|
||
|
|
self.document_border = self.get_document_borders_as_beamdump()
|
||
|
|
self.world.add(self.document_border)
|
||
|
|
|
||
|
|
filter_ = self.filter_primitives + (inkex.Group, inkex.Use)
|
||
|
|
for obj in self.svg.selection.filter(filter_):
|
||
|
|
self.add(obj)
|
||
|
|
|
||
|
|
if self.beam_seeds:
|
||
|
|
for seed in self.beam_seeds:
|
||
|
|
if self.is_inside_document(seed.ray):
|
||
|
|
generated = self.world.propagate_beams(seed.ray)
|
||
|
|
for beam in generated:
|
||
|
|
try:
|
||
|
|
new_layer = get_or_create_beam_layer(
|
||
|
|
get_containing_layer(seed.parent)
|
||
|
|
)
|
||
|
|
plot_beam(beam, seed.parent, new_layer)
|
||
|
|
except LayerError as e:
|
||
|
|
inkex.utils.errormsg(f"{e} It will be ignored.")
|
||
|
|
|
||
|
|
@singledispatchmethod
|
||
|
|
def add(self, obj):
|
||
|
|
pass
|
||
|
|
|
||
|
|
@add.register
|
||
|
|
def _(self, group: inkex.Group):
|
||
|
|
for child in group:
|
||
|
|
self.add(child)
|
||
|
|
|
||
|
|
@add.register
|
||
|
|
def _(self, clone: inkex.Use):
|
||
|
|
copy = get_unlinked_copy(clone)
|
||
|
|
self.add(copy)
|
||
|
|
|
||
|
|
for type in filter_primitives:
|
||
|
|
|
||
|
|
@add.register(type)
|
||
|
|
def _(self, obj):
|
||
|
|
"""
|
||
|
|
Extracts properties and adds the object to the ray tracing data
|
||
|
|
structure
|
||
|
|
"""
|
||
|
|
material = get_material(obj)
|
||
|
|
if material:
|
||
|
|
if isinstance(material, BeamSeed):
|
||
|
|
for ray in get_beams(obj):
|
||
|
|
self.beam_seeds.append(BeamSeed(ray, parent=obj))
|
||
|
|
else:
|
||
|
|
geometry = get_geometry(obj)
|
||
|
|
opt_obj = OpticalObject(geometry, material)
|
||
|
|
self.world.add(opt_obj)
|
||
|
|
|
||
|
|
def get_document_borders_as_beamdump(self) -> OpticalObject:
|
||
|
|
"""
|
||
|
|
Adds a beam blocking contour on the borders of the document to
|
||
|
|
prevent the beams from going to infinity
|
||
|
|
"""
|
||
|
|
|
||
|
|
c1x, c1y, c2x, c2y = self.svg.get_viewbox()
|
||
|
|
contour_geometry = CompoundGeometricObject(
|
||
|
|
(
|
||
|
|
CubicBezier(
|
||
|
|
Vector(c1x, c1y),
|
||
|
|
Vector(c1x, c1y),
|
||
|
|
Vector(c2x, c1y),
|
||
|
|
Vector(c2x, c1y),
|
||
|
|
),
|
||
|
|
CubicBezier(
|
||
|
|
Vector(c2x, c1y),
|
||
|
|
Vector(c2x, c1y),
|
||
|
|
Vector(c2x, c2y),
|
||
|
|
Vector(c2x, c2y),
|
||
|
|
),
|
||
|
|
CubicBezier(
|
||
|
|
Vector(c2x, c2y),
|
||
|
|
Vector(c2x, c2y),
|
||
|
|
Vector(c1x, c2y),
|
||
|
|
Vector(c1x, c2y),
|
||
|
|
),
|
||
|
|
CubicBezier(
|
||
|
|
Vector(c1x, c2y),
|
||
|
|
Vector(c1x, c2y),
|
||
|
|
Vector(c1x, c1y),
|
||
|
|
Vector(c1x, c1y),
|
||
|
|
),
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return OpticalObject(contour_geometry, raytracing.material.BeamDump())
|
||
|
|
|
||
|
|
def is_inside_document(self, ray: Ray) -> bool:
|
||
|
|
return self.document_border.geometry.is_inside(ray)
|
||
|
|
|
||
|
|
|
||
|
|
def get_material(
|
||
|
|
obj: inkex.ShapeElement,
|
||
|
|
) -> Optional[raytracing.material.OpticMaterial | BeamSeed]:
|
||
|
|
"""Extracts the optical material of an object from its description"""
|
||
|
|
|
||
|
|
desc = obj.desc
|
||
|
|
if desc is None:
|
||
|
|
desc = ""
|
||
|
|
materials = get_materials_from_description(desc)
|
||
|
|
if len(materials) == 0:
|
||
|
|
return None
|
||
|
|
if len(materials) > 1:
|
||
|
|
raise_err_num_materials(obj)
|
||
|
|
elif len(materials) == 1:
|
||
|
|
return materials[0]
|
||
|
|
|
||
|
|
|
||
|
|
def get_materials_from_description(
|
||
|
|
desc: str,
|
||
|
|
) -> list[raytracing.material.OpticMaterial | BeamSeed]:
|
||
|
|
"""Run through the description to extract the material properties"""
|
||
|
|
|
||
|
|
materials = list()
|
||
|
|
class_alias = dict(
|
||
|
|
beam_dump=raytracing.material.BeamDump,
|
||
|
|
mirror=raytracing.material.Mirror,
|
||
|
|
beam_splitter=raytracing.material.BeamSplitter,
|
||
|
|
glass=raytracing.material.Glass,
|
||
|
|
beam=BeamSeed,
|
||
|
|
)
|
||
|
|
for match in get_optics_fields(desc):
|
||
|
|
material_type = match.group("material")
|
||
|
|
prop_str = match.group("num")
|
||
|
|
if material_type in class_alias:
|
||
|
|
if material_type == "glass" and prop_str is not None:
|
||
|
|
optical_index = float(prop_str)
|
||
|
|
materials.append(class_alias[material_type](optical_index))
|
||
|
|
else:
|
||
|
|
materials.append(class_alias[material_type]())
|
||
|
|
return materials
|
||
|
|
|
||
|
|
|
||
|
|
def raise_err_num_materials(obj):
|
||
|
|
inkex.utils.errormsg(
|
||
|
|
f"The element {obj.get_id()} has more than one optical material and will be"
|
||
|
|
f" ignored:\n{obj.desc}\n"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def get_geometry(obj: inkex.ShapeElement) -> GeometricObject:
|
||
|
|
"""
|
||
|
|
Converts the geometry of inkscape elements to a form suitable for the
|
||
|
|
ray tracing module
|
||
|
|
"""
|
||
|
|
|
||
|
|
# Treats all objects as cubic Bezier curves. This treatment is exact
|
||
|
|
# for most primitives except circles and ellipses that are only
|
||
|
|
# approximated by Bezier curves.
|
||
|
|
# TODO: implement exact representation for ellipses
|
||
|
|
path = get_absolute_path(obj)
|
||
|
|
composite_bezier = convert_to_composite_bezier(path)
|
||
|
|
return composite_bezier
|
||
|
|
|
||
|
|
|
||
|
|
def get_absolute_path(obj: inkex.ShapeElement) -> inkex.CubicSuperPath:
|
||
|
|
path = obj.to_path_element().path.to_absolute()
|
||
|
|
transformed_path = path.transform(obj.composed_transform())
|
||
|
|
return transformed_path.to_superpath()
|
||
|
|
|
||
|
|
|
||
|
|
def get_beams(element: inkex.ShapeElement) -> Iterable[Ray]:
|
||
|
|
"""
|
||
|
|
Returns a beam with origin at the endpoint of the path and tangent to
|
||
|
|
the path
|
||
|
|
"""
|
||
|
|
bezier_path = convert_to_composite_bezier(get_absolute_path(element))
|
||
|
|
for sub_path in bezier_path:
|
||
|
|
last_segment = sub_path[-1]
|
||
|
|
endpoint = last_segment.eval(1)
|
||
|
|
tangent = -last_segment.tangent(1)
|
||
|
|
yield Ray(endpoint, tangent)
|
||
|
|
|
||
|
|
|
||
|
|
def convert_to_composite_bezier(
|
||
|
|
superpath: inkex.CubicSuperPath,
|
||
|
|
) -> CompoundGeometricObject:
|
||
|
|
"""
|
||
|
|
Converts a superpath with a representation
|
||
|
|
[Subpath0[handle0_0, point0, handle0_1], ...], ...]
|
||
|
|
to a representation of consecutive bezier segments of the form
|
||
|
|
CompositeCubicBezier([CubicBezierPath[CubicBezier[point0, handle0_1,
|
||
|
|
handle1_0, point1], ...], ...]).
|
||
|
|
"""
|
||
|
|
|
||
|
|
composite_bezier = list()
|
||
|
|
for subpath in superpath:
|
||
|
|
bezier_path = list()
|
||
|
|
for (__, p0, p1), (p2, p3, __) in pairwise(subpath):
|
||
|
|
bezier = CubicBezier(Vector(*p0), Vector(*p1), Vector(*p2), Vector(*p3))
|
||
|
|
bezier_path.append(bezier)
|
||
|
|
composite_bezier.append(CompoundGeometricObject(bezier_path))
|
||
|
|
return CompoundGeometricObject(composite_bezier)
|
||
|
|
|
||
|
|
|
||
|
|
def get_containing_layer(obj: inkex.BaseElement) -> inkex.Layer:
|
||
|
|
try:
|
||
|
|
return obj.ancestors().filter(inkex.Layer)[0]
|
||
|
|
except IndexError:
|
||
|
|
raise LayerError(f"Object '{obj.get_id()}' is not inside a layer.")
|
||
|
|
|
||
|
|
|
||
|
|
class LayerError(RuntimeError):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
Raytracing().run()
|