Skip to content

Commit

Permalink
add doc strings, fix test name
Browse files Browse the repository at this point in the history
  • Loading branch information
@chrisjonesBSU
chrisjonesBSU committed Sep 18, 2023
1 parent 654ab42 commit 2a22909
Show file tree
Hide file tree
Showing 2 changed files with 83 additions and 13 deletions.
5 changes: 1 addition & 4 deletions cmeutils/tests/test_visualize.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,9 @@ class TestFresnelGSD(BaseTest):
def test_view(self, p3ht_fresnel):
p3ht_fresnel.view()

def test_path_frace(self, p3ht_fresnel):
def test_path_trace(self, p3ht_fresnel):
p3ht_fresnel.path_trace(samples=10, light_samples=1)

def test_frace(self, p3ht_fresnel):
p3ht_fresnel.trace()

def test_scale_diameter(self, p3ht_fresnel):
p3ht_fresnel.diameter_scale = 0.6
assert p3ht_fresnel.diameter_scale == 0.6
Expand Down
91 changes: 82 additions & 9 deletions cmeutils/visualize.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ def __init__(
frame=0,
color_dict=None,
diameter_scale=0.30,
solid=0.1,
solid=0,
roughness=0.3,
specular=0.5,
specular_trans=0,
Expand Down Expand Up @@ -43,6 +43,9 @@ def __init__(

@property
def frame(self):
"""The frame of the GSD trajectory to use.
The frame determines the particle positions usedin the image.
"""
return self._frame

@frame.setter
Expand All @@ -55,10 +58,14 @@ def frame(self, frame):

@property
def snapshot(self):
"""gsd.hoomd.Snapshot loaded from the GSD file.
The snapshot loaded depends on FresnelGSD.frame
"""
return self._snapshot

@property
def diameter_scale(self):
"""Scales the snapshot diameter values to set FresnelGSD.radius"""
return self._diameter_scale

@diameter_scale.setter
Expand All @@ -67,6 +74,7 @@ def diameter_scale(self, value):

@property
def color_dict(self):
"""Dictionary of key: particle type value: color"""
return self._color_dict

@color_dict.setter
Expand All @@ -79,6 +87,16 @@ def color_dict(self, value):
self._color_dict = value

def set_type_color(self, particle_type, color):
"""Set colors for particle types one at a time
Parameters
----------
particle_type : str; required
The particle type found in FresnelGSD.particle_types
or FresnelGSD.snapshot.particles.types
color : np.ndarray, shape=(3,), required
sRGB color of (red, green, blue) values
"""
if particle_type not in set(self.particle_types):
raise ValueError(
f"Particle type of {particle_type} is not in the Snapshot"
Expand All @@ -87,6 +105,12 @@ def set_type_color(self, particle_type, color):

@property
def unwrap_positions(self):
"""If set to True, then positions of the particles are
unwrapped before creating the image.
This requires that the GSD file snapshots contain accurate
image values (gsd.hoomd.Snapshot.particles.image
"""
return self._unwrap_positions

@unwrap_positions.setter
Expand All @@ -100,6 +124,7 @@ def unwrap_positions(self, value):

@property
def solid(self):
"""Set to 1 to use solid colors regardless of light and angle"""
return self._solid

@solid.setter
Expand All @@ -108,6 +133,7 @@ def solid(self, value):

@property
def roughness(self):
"""Sets the material roughness (ranges from 0.1 to 1)"""
return self._roughness

@roughness.setter
Expand All @@ -116,6 +142,9 @@ def roughness(self, value):

@property
def specular(self):
"""Determines the strength of the specular highlights
(ranges from 0 to 1)
"""
return self._specular

@specular.setter
Expand All @@ -124,6 +153,9 @@ def specular(self, value):

@property
def specular_trans(self):
"""Determines the magnitude of specular light transmission
(ranges from 0 to 1)
"""
return self._specular_trans

@specular_trans.setter
Expand All @@ -132,6 +164,9 @@ def specular_trans(self, value):

@property
def metal(self):
"""Set to 0 for dielectric material or 1 for metal
(ranges from 0 to 1)
"""
return self._metal

@metal.setter
Expand All @@ -140,6 +175,7 @@ def metal(self, value):

@property
def view_axis(self):
"""Sets the direction and position of the camera"""
return self._view_axis

@view_axis.setter
Expand All @@ -149,6 +185,7 @@ def view_axis(self, value):

@property
def up(self):
"""Determines which direction is up"""
return self._up

@up.setter
Expand All @@ -157,6 +194,7 @@ def up(self, value):

@property
def height(self):
"""Acts like a zoom. Larger values zoom out, smaller vaues zoom in"""
return self._height

@height.setter
Expand All @@ -165,14 +203,22 @@ def height(self, value):

@property
def camera_position(self):
"""The camera position.
Determined from box dimensions are FresnelGSD.view_axis"""
return (self.snapshot.configuration.box[:3] * self.view_axis) - 0.5

@property
def look_at(self):
"""The direction the camera is facing.
By default, uses position directly opposite of camera position
"""
return self.snapshot.configuration.box[:3] * -self.view_axis

@property
def positions(self):
"""Particle positions used in the image.
Determined by FresnelGSD.frame
"""
if self.unwrap_positions:
pos = self.snapshot.particles.position
imgs = self.snapshot.particles.image
Expand All @@ -183,10 +229,15 @@ def positions(self):

@property
def radius(self):
"""Sets the size of the particles.
Determined by the gsd.hoomd.Snapshot.particles.diameter
values and FresnelGSD.diameter_scale
"""
return self.snapshot.particles.diameter * self.diameter_scale

@property
def particle_types(self):
"""Array of particle types of length number of particles"""
return np.array(
[
self.snapshot.particles.types[i]
Expand All @@ -196,12 +247,16 @@ def particle_types(self):

@property
def colors(self):
"""Generates the colors by particle type,
or sets a default of (0.5, 0.25, 0.5)
"""
if self.color_dict:
return np.array([self.color_dict[i] for i in self.particle_types])
else:
return np.array([0.5, 0.25, 0.5])

def geometry(self):
"""Creates and returns a fresnel.geometry.Sphere object"""
geometry = fresnel.geometry.Sphere(
scene=self.scene,
position=self.positions,
Expand All @@ -212,6 +267,7 @@ def geometry(self):
return geometry

def material(self):
"""Creates and returns a fresnel.material.Material object"""
material = fresnel.material.Material(
primitive_color_mix=1,
solid=self.solid,
Expand All @@ -223,6 +279,7 @@ def material(self):
return material

def camera(self):
"""Creates and returns a fresnel.camera.Orthographic object"""
camera = fresnel.camera.Orthographic(
position=self.camera_position,
look_at=self.look_at,
Expand All @@ -232,11 +289,35 @@ def camera(self):
return camera

def view(self, width=300, height=300):
"""Creates an image using fresnel.preview
Parameters
----------
width : int, optional, default 300
Image width size in pixels
height : int, optional, default 300
Image height size in pixels
"""
self.scene.camera = self.camera()
self.scene.geometry = [self.geometry()]
return fresnel.preview(scene=self.scene, w=width, h=height)

def path_trace(self, width=300, height=300, samples=64, light_samples=1):
"""Creates an image using fresnel.pathtrace
Parameters
----------
width : int, optional, default 300
Image width size in pixels
height : int, optional, default 300
Image height size in pixels
samples : int, optional, default 64
The number of times to sample pixels in the scene
light_samples : int, optional, default=1
The number of light samples for each pixel sample
"""
self.scene.camera = self.camera()
self.scene.geometry = [self.geometry()]
return fresnel.pathtrace(
Expand All @@ -246,11 +327,3 @@ def path_trace(self, width=300, height=300, samples=64, light_samples=1):
samples=samples,
light_samples=light_samples,
)

def trace(self, width=300, height=300, n_samples=1):
self.scene.camera = self.camera()
self.scene.geometry = [self.geometry()]
tracer = fresnel.tracer.Preview(
device=self.scene.device, w=width, h=height
)
return tracer.render(self.scene)

0 comments on commit 2a22909

Please sign in to comment.