some minor enhancements

Author: Fe-r-oz

docs/src/users_guide/plotting_the_bloch_sphere.md

@@ -121,9 +121,9 @@ fig, ax = render(b);
 fig
 ```
 
-## Add Another Set of Points
+Notice that, in contrast to states or vectors, each point remains the same color as the initial point. This is because adding multiple data points using [`add_points!`](@ref) is interpreted, by default, to correspond to a single data point (single qubit state) plotted at different times. This is very useful when visualizing the dynamics of a qubit. An example of this is given in the example . If we want to plot additional qubit states we can call additional [`add_points!`](@ref) functions:
 
-If we want to plot additional qubit states we can call additional [`add_points!`](@ref) functions:
+## Add Another Set of Points
 
 ```@example Bloch_sphere_rendering
 xz = zeros(20);
@@ -134,3 +134,27 @@ add_points!(b, pnts);
 fig, ax = render(b);
 fig
 ```
+
+The color and shape of the data points is varied automatically by [`Bloch`](@ref). Notice how the color and point markers change for each set of data. Again, we have had to call add_points twice because adding more than one set of multiple data points is not supported by the add_points function.
+
+What if we want to vary the color of our points. We can tell [`Bloch`](@ref) to vary the color of each point according to the colors listed in the `b.point_color`
+
+```@example Bloch_sphere_rendering
+clear!(b)
+xp = cos.(th);
+yp = sin.(th);
+zp = zeros(20);
+pnts = [xp, yp, zp];
+add_points!(b, pnts, meth=:m);
+fig, ax = render(b);
+fig
+```
+
+Now, the data points cycle through a variety of predefined colors. Now lets add another set of points, but this time we want the set to be a single color, representing say a qubit going from the ``|0\rangle`` state to the state in the `y-z` plane:
+
+```@example Bloch_sphere_rendering
+pnts = [xz, yz, zz] ;
+add_points!(b, pnts);
+fig, ax = render(b);
+fig
+```

ext/QuantumToolboxMakieExt.jl

@@ -630,27 +630,62 @@ Plot all quantum state points on the Bloch sphere.
 Handles both scatter points and line traces based on style specifications.
 """
 function _plot_points!(b::Bloch, ax)
-    for (k, points) in enumerate(b.points)
+    for k in 1:length(b.points)
+        pts = b.points[k]
         style = b.point_style[k]
-        color = b.point_color[k]
         alpha = b.point_alpha[k]
-        marker = b.point_marker[(k-1)%length(b.point_marker)+1]
-        x, y, z = points[2, :], -points[1, :], points[3, :]
-        if style ∈ (:s, :m)
+        marker = b.point_marker[mod1(k, length(b.point_marker))]
+        N = size(pts, 2)
+
+        raw_x = pts[2, :]
+        raw_y = -pts[1, :]
+        raw_z = pts[3, :]
+
+        ds = vec(sqrt.(sum(abs2, pts; dims = 1)))
+        if !all(isapprox.(ds, ds[1]; rtol = 1e-12))
+            indperm = sortperm(ds)
+        else
+            indperm = collect(1:N)
+        end
+        this_color = b.point_color[k]
+        if style == :m
+            defaults = b.point_default_color
+            L = length(defaults)
+            times = ceil(Int, N / L)
+            big_colors = repeat(b.point_default_color, times)[1:N]
+            big_colors = big_colors[indperm]
+            colors = big_colors
+        else
+            if this_color === nothing
+                defaults = b.point_default_color
+                colors = defaults[mod1(k, length(defaults))]
+            else
+                colors = this_color
+            end
+        end
+        if style in (:s, :m)
+            xs = raw_x[indperm]
+            ys = raw_y[indperm]
+            zs = raw_z[indperm]
             scatter!(
                 ax,
-                x,
-                y,
-                z;
-                color = color,
-                markersize = 6,
+                xs,
+                ys,
+                zs;
+                color = colors,
+                markersize = b.point_size[mod1(k, length(b.point_size))],
                 marker = marker,
-                strokewidth = 0.05,
-                transparency = alpha<1,
+                transparency = alpha < 1.0,
                 alpha = alpha,
+                strokewidth = 0.0,
             )
+
         elseif style == :l
-            lines!(ax, x, y, z; color = color, linewidth = 2.0, transparency = alpha<1, alpha = alpha)
+            xs = raw_x
+            ys = raw_y
+            zs = raw_z
+            c = isa(colors, Vector) ? colors[1] : colors
+            lines!(ax, xs, ys, zs; color = c, linewidth = 2.0, transparency = alpha < 1.0, alpha = alpha)
         end
     end
 end

src/visualization.jl

@@ -137,10 +137,10 @@ A structure representing a Bloch sphere visualization for quantum states.
     frame_alpha::Float64 = 0.0
     frame_color::String = "white"
     frame_limit::Float64 = 1.5
-    point_default_color::Vector{String} = ["blue", "red", "green", "orange", "cyan", "magenta", "yellow", "black"]
-    point_color::Vector{String} = ["blue", "red", "green", "orange", "cyan", "magenta", "yellow", "black"]
+    point_default_color::Vector{String} = ["blue", "red", "green", "orange"]
+    point_color::Vector{Union{Nothing,String}} = Union{Nothing,String}[]
     point_marker::Vector{Symbol} = [:circle, :rect, :diamond, :utriangle]
-    point_size::Vector{Float64} = [4.0, 4.8, 5.0, 6.0]
+    point_size::Vector{Float64} = [5.0, 6.0, 7.0, 8.0]
     point_style::Vector{Symbol} = Symbol[]
     point_alpha::Vector{Float64} = Float64[]
     sphere_alpha::Float64 = 0.4
@@ -228,20 +228,28 @@ Add multiple points to the Bloch sphere visualization.
 - alpha=1.0: Transparency (1.0 = opaque, 0.0 = transparent)
 ```
 """
-function add_points!(b::Bloch, pnts::Matrix{<:Real}; meth::Symbol = :s, color = nothing, alpha = 1.0)
+function add_points!(
+    b::Bloch,
+    pnts::Matrix{<:Real};
+    meth::Symbol = :s,
+    color::Union{Nothing,String} = nothing,
+    alpha::Float64 = 1.0,
+)
     if size(pnts, 1) != 3
-        error("Points must be a 3×N matrix where columns are [x;y;z] points")
+        error("Points must be a 3×N matrix where each column is [x; y; z]")
     end
-    if meth ∉ [:s, :m, :l]
-        error("meth must be :s, :m, or :l")
-    end
-    if meth == :s && size(pnts, 2) == 1
-        pnts = hcat(pnts, pnts)
+    if !(meth in (:s, :m, :l))
+        error("`meth` must be :s, :m, or :l")
     end
     push!(b.points, convert(Matrix{Float64}, pnts))
     push!(b.point_style, meth)
     push!(b.point_alpha, alpha)
-    return push!(b.point_color, color === nothing ? b.point_default_color[1] : color)
+    if color === nothing
+        push!(b.point_color, nothing)
+    else
+        push!(b.point_color, color)
+    end
+    return nothing
 end
 
 @doc raw"""
@@ -324,6 +332,9 @@ Clear all graphical elements (points, vectors, lines, arcs) from the given Bloch
 """
 function clear!(b::Bloch)
     empty!(b.points)
+    empty!(b.point_color)
+    empty!(b.point_style)
+    empty!(b.point_alpha)
     empty!(b.vectors)
     empty!(b.lines)
     empty!(b.arcs)