From d14d296274cccde385302b6cc77371430d7af257 Mon Sep 17 00:00:00 2001
From: Tao Wang <fsxbhyy@gmail.com>
Date: Fri, 4 Nov 2022 11:50:10 -0400
Subject: [PATCH] add 2D irreducible plot

---
 toolbox/plotlyjs_wignerseitz.jl | 21 ++++++++++++++++++++-
 toolbox/test_extend.jl          | 29 +++++++++++++++++------------
 2 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/toolbox/plotlyjs_wignerseitz.jl b/toolbox/plotlyjs_wignerseitz.jl
index a066798..acee13a 100644
--- a/toolbox/plotlyjs_wignerseitz.jl
+++ b/toolbox/plotlyjs_wignerseitz.jl
@@ -364,7 +364,7 @@ function plot(c::Cell{2}, layout::Layout=Layout();
     return PlotlyJS.plot(ts, layout; config=config)
 end
 
-function plot(clist::Array{Cell{2}}, idx_center::Int, layout::Layout=Layout();
+function plot(clist::Array{Cell{2}}, idx_center::Int, layout::Layout=Layout(); ibz =nothing,
     config::PlotConfig=PlotConfig(responsive=true, displaylogo=false))
 
     layout = merge(DEFAULT_PLOTLY_LAYOUT_2D, layout)
@@ -380,6 +380,25 @@ function plot(clist::Array{Cell{2}}, idx_center::Int, layout::Layout=Layout();
     max_x, max_y = maximum(v -> abs(v[1]), basis(c_center)), maximum(v -> abs(v[2]), basis(c_center))
     get!(layout[:xaxis], :range, [-max_x - scale / 15, max_x + scale / 15])
     get!(layout[:yaxis], :range, [-max_y - scale / 15, max_y + scale / 15])
+
+
+    if !isnothing(ibz)
+        setting(ibz) !== CARTESIAN && (ibz = cartesianize(ibz))        
+        # merge_coplanar!(ibz) # Coplanar triangles have to be merged after calling mesh3d.
+        # BZ
+        merge_coplanar!(ibz)
+        tbz = Vector{GenericTrace{Dict{Symbol,Any}}}(undef, length(ibz))
+        for (i, poly) in enumerate(ibz)
+            tbz[i] = PlotlyJS.scatter(
+                x=push!(getindex.(poly, 1), poly[1][1]),
+                y=push!(getindex.(poly, 2), poly[1][2]);
+                mode="lines", hovertext="Cell", hoverinfo="text+x+y",
+                line=attr(color=BZ_COL[], width=3 ), fill = "toself", fillcolor = "rgb(50, 200, 200)"
+            )
+        end
+        push!(tbz_list, tbz)
+    end
+
     for (ci, c) in enumerate(clist)
         if (!isnothing(c))
             # Cell boundaries
diff --git a/toolbox/test_extend.jl b/toolbox/test_extend.jl
index b04a59b..6225219 100644
--- a/toolbox/test_extend.jl
+++ b/toolbox/test_extend.jl
@@ -106,23 +106,22 @@ end
 
 # Wigner-Seitz cells visualization
 # 2D
-#Rs = [[1.0, 0.0], [-0.5, √3/2]]
 #N, DIM = 4, 2
 #origin = [0.0 0.0]
-#lattice =Matrix([2 0; 1 sqrt(3)]')
-#msize = (3,3)
+lattice =Matrix([2 0; 1 sqrt(3)]')
 #br = BZMeshes.Brillouin(lattice = lattice)
 
 # latvec = [1 0; 0 1]'
 # 3D
 
-lattice = [[0 0.5 0.5]; [0.5 0 0.5]; [0.5 0.5 0.0]]
+#lattice = [[0 0.5 0.5]; [0.5 0 0.5]; [0.5 0.5 0.0]]
 #atoms = [1,1]
 #positions = [ones(3) / 8, -ones(3) / 8]
 
 #lattice = [[1.0 0.0 0.0]; [0.0 1.0 0.0]; [0.0 0.0 1.0]]
 atoms = [1]
-positions = [zeros(3)]
+#positions = [zeros(3)]
+positions = [zeros(2)]
 
 atom_pos = hvcat(size(positions, 1), positions...)
 ibzformat = "convex hull"
@@ -132,13 +131,16 @@ convention = "ordinary"
 
 br = BZMeshes.Brillouin(lattice=lattice, atoms=atoms, positions=positions)
 #br = BZMeshes.Brillouin(lattice=lattice)
-msize = (8, 8, 8)
 
-bzmesh = UniformBZMesh(br=br, size=msize, shift=[false, false, false])
-meshmap = MeshMap(bzmesh)
+# msize = (8, 8, 8)
+# bzmesh = UniformBZMesh(br=br, size=msize, shift=[false, false, false])
+
+msize = (3,3)
+bzmesh = UniformBZMesh(br=br, size=msize, shift=[false, false])
+
 
+meshmap = MeshMap(bzmesh)
 recip_lattice = lattice_vector(bzmesh)
-#latvec = mapslices(x -> [x],recip_lattice , dims=1)[:]
 latvec = mapslices(x -> [x], recip_lattice, dims=1)[:]
 
 #generate irreducible brillouin zone
@@ -179,8 +181,11 @@ end
 
 P = plot(clist, idx_center, ibz=c)
 
-addtraces!(P, scatter3d(x=[r[1] for r in fullmesh], y=[r[2] for r in fullmesh], z=[r[3] for r in fullmesh], mode="markers", marker=attr(size=3)))
-addtraces!(P, scatter3d(x=[r[1] for r in reducedmesh], y=[r[2] for r in reducedmesh], z=[r[3] for r in reducedmesh], mode="markers", marker=attr(size=3)))
+addtraces!(P, scatter(x=[r[1] for r in fullmesh], y=[r[2] for r in fullmesh], mode="markers", marker=attr(size=3)))
+addtraces!(P, scatter(x=[r[1] for r in reducedmesh], y=[r[2] for r in reducedmesh], mode="markers", marker=attr(size=3)))
+
+#addtraces!(P, scatter3d(x=[r[1] for r in fullmesh], y=[r[2] for r in fullmesh], z=[r[3] for r in fullmesh], mode="markers", marker=attr(size=3)))
+#addtraces!(P, scatter3d(x=[r[1] for r in reducedmesh], y=[r[2] for r in reducedmesh], z=[r[3] for r in reducedmesh], mode="markers", marker=attr(size=3)))
 
 n = length(P.plot.data) # number of traces, the last two are the full and reduced meshes
 allinvis, vis1, vis2 = [true for i in 1:n], [true for i in 1:n], [true for i in 1:n]
@@ -231,5 +236,5 @@ d = Dict(
 relayout!(P, d)
 
 display(P)
-# readline()
+readline()