From ff06985913c5f33e663361eb07cb8f4c46fc8ecb Mon Sep 17 00:00:00 2001
From: Sikorski <sikorski@zib.de>
Date: Tue, 13 Aug 2024 16:50:00 +0200
Subject: [PATCH] IsoSimulation: pretty print pairwise features

---
 src/simulation.jl        | 5 ++++-
 src/simulators/openmm.jl | 8 +++++++-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/simulation.jl b/src/simulation.jl
index 8060881..6865368 100644
--- a/src/simulation.jl
+++ b/src/simulation.jl
@@ -180,10 +180,13 @@ function resample_kde(data, model, n; padding=0.0, bandwidth=0.02)
 end
 
 function Base.show(io::IO, mime::MIME"text/plain", d::SimulationData)#
+    simstr = sprint() do io
+        show(io, mime, d.sim)
+    end
     println(
         io, """
         SimulationData(;
-            sim=$(d.sim),
+            sim=$(simstr),
             features=$(size.(d.features)), $(split(string(typeof(d.features[1])),",")[1]),
             coords=$(size.(d.coords)), $(split(string(typeof(d.coords[1])),",")[1]),
             featurizer=$(d.featurizer))"""
diff --git a/src/simulators/openmm.jl b/src/simulators/openmm.jl
index 3216ee3..a28fe15 100644
--- a/src/simulators/openmm.jl
+++ b/src/simulators/openmm.jl
@@ -354,6 +354,11 @@ Base.convert(::Type{OpenMMSimulation}, a::OpenMMSimulationSerialized) =
         mmthreads=a.mmthreads)
 
 function Base.show(io::IO, mime::MIME"text/plain", sim::OpenMMSimulation)#
+    featstr = if sim.features isa Vector{Tuple{Int,Int}}
+        "{$(length(sim.features)) pairwise distances}"
+    else
+        sim.features
+    end
     println(
         io, """
         OpenMMSimulation(;
@@ -364,7 +369,8 @@ function Base.show(io::IO, mime::MIME"text/plain", sim::OpenMMSimulation)#
             friction=$(sim.friction),
             step=$(sim.step),
             steps=$(sim.steps),
-            features=$(sim.features))"""
+            features=$featstr) 
+        with $(div(length(getcoords(sim)),3)) atoms"""
     )
 end