1package main
  2
  3import (
  4	"image/color"
  5	"log"
  6	"math"
  7	"time"
  8
  9	"gonum.org/v1/gonum/mat"
 10
 11	"github.com/alltom/oklab"
 12	"github.com/hajimehoshi/ebiten/v2"
 13	"github.com/hajimehoshi/ebiten/v2/inpututil"
 14	"github.com/hajimehoshi/ebiten/v2/vector"
 15)
 16
 17const (
 18	N                       = 100
 19	L                       = 1e-9            // 1 nm box
 20	hbar                    = 1.054571817e-34 // J·s
 21	mass                    = 9.10938356e-31  // kg
 22	eV                      = 1.602176634e-19 // J
 23	V0Max                   = 30 * eV
 24	ProbabilityDensityScale = 0.3       // Percentage of the screen space
 25	TimeFactor              = 2 * 1e-25 // Time factor for animation
 26	Width                   = 1200
 27	Height                  = 1200
 28)
 29
 30// Game implements ebiten.Game interface.
 31type Game struct {
 32	changedPotential bool
 33
 34	potential              []float64
 35	eigenvalues            []float64
 36	eigenvectors           mat.Dense
 37	maxdisplayedeigenvalue int
 38	computedColors         []color.RGBA
 39
 40	renormalizedEigenvectors mat.Dense
 41
 42	mousePressed bool
 43	prevX        int
 44	prevY        float64
 45
 46	touchIDs         []ebiten.TouchID
 47	releasedTouchIDs []ebiten.TouchID
 48}
 49
 50// Update proceeds the game state.
 51// Update is called every tick (1/60 [s] by default).
 52func (g *Game) Update() error {
 53	// Write your game's logical update.
 54	if g.changedPotential {
 55		dx := L / float64(N)
 56
 57		H := mat.NewSymDense(N, nil)
 58
 59		// kinetic prefactor
 60		t := hbar * hbar / (2 * mass * dx * dx)
 61
 62		for i := range N {
 63			V := g.potential[i] * V0Max
 64
 65			H.SetSym(i, i, 2*t+V)
 66
 67			if i > 0 {
 68				H.SetSym(i, i-1, -t)
 69			}
 70		}
 71
 72		var eig mat.EigenSym
 73		ok := eig.Factorize(H, true)
 74		if !ok {
 75			panic("eigendecomposition failed")
 76		}
 77		g.eigenvalues = eig.Values(nil)
 78		eig.VectorsTo(&g.eigenvectors)
 79
 80		g.maxdisplayedeigenvalue = -1
 81		g.computedColors = []color.RGBA{}
 82		for i := range g.eigenvalues {
 83			if g.eigenvalues[i] <= V0Max {
 84				g.maxdisplayedeigenvalue = i
 85				g.computedColors = append(g.computedColors, OKLCHToColor(1, 0.5, 2*math.Pi*g.eigenvalues[i]/V0Max))
 86				norm := 0.0
 87				for j := range N {
 88					norm += math.Pow(g.eigenvectors.At(j, i), 2)
 89				}
 90				norm = math.Sqrt(norm) / ProbabilityDensityScale
 91				for j := range N {
 92					y := float64(Height) * float64(g.eigenvectors.At(j, i)/norm)
 93					g.eigenvectors.Set(j, i, y)
 94				}
 95			}
 96		}
 97
 98		g.changedPotential = false
 99	}
100
101	g.touchIDs = ebiten.AppendTouchIDs(g.touchIDs[:0])
102
103	if ebiten.IsMouseButtonPressed(ebiten.MouseButtonLeft) || len(g.touchIDs) != 0 {
104		x, y := ebiten.CursorPosition()
105		if len(g.touchIDs) != 0 {
106			x, y = ebiten.TouchPosition(g.touchIDs[0])
107		}
108		newX := int(float64(x) / float64(Width) * (N - 1))
109		newY := float64(y) / float64(Height)
110		if g.mousePressed {
111			diff := newX - g.prevX
112			if diff > 0 {
113				for i := range diff {
114					j := i + g.prevX
115					if j >= 0 && j < N {
116						g.potential[j] = 1 - (newY*float64(i)/float64(diff) + g.prevY*(1-float64(i)/float64(diff)))
117					}
118				}
119			} else {
120				for i := range -(diff) {
121					j := g.prevX - i - 1
122					if j >= 0 && j < N {
123						g.potential[j] = 1 - (newY*float64(i)/float64(-diff) + g.prevY*(1-float64(i)/float64(-diff)))
124					}
125				}
126			}
127		}
128		g.prevX = newX
129		g.prevY = newY
130		g.mousePressed = true
131
132	}
133
134	g.releasedTouchIDs = inpututil.AppendJustReleasedTouchIDs(g.releasedTouchIDs[:0])
135
136	if inpututil.IsMouseButtonJustReleased(ebiten.MouseButtonLeft) || len(g.releasedTouchIDs) != 0 {
137		g.changedPotential = true
138		g.mousePressed = false
139	}
140
141	return nil
142}
143
144func OKLCHToColor(L, C, H float64) color.RGBA {
145	c := oklab.Oklch{
146		L: L,
147		C: C,
148		H: H,
149	}
150
151	r, g, b, a := c.RGBA()
152
153	return color.RGBA{
154		R: uint8(r >> 8),
155		G: uint8(g >> 8),
156		B: uint8(b >> 8),
157		A: uint8(a >> 8),
158	}
159}
160
161// Draw draws the game screen.
162// Draw is called every frame (typically 1/60[s] for 60Hz display).
163func (g *Game) Draw(screen *ebiten.Image) {
164	screen.Fill(color.RGBA{0, 0, 0, 255})
165	for i := range N - 1 {
166		vector.StrokeLine(screen,
167			float32(Width)*(float32(i)/float32(N-1)),
168			float32(Height)-float32(Height)*float32(g.potential[i]),
169			float32(Width)*(float32(i+1)/float32(N-1)),
170			float32(Height)-float32(Height)*float32(g.potential[i+1]),
171			2,
172			color.RGBA{255, 255, 255, 255},
173			true)
174	}
175	for i := range g.maxdisplayedeigenvalue {
176		if g.eigenvalues[i] <= V0Max {
177			phase := math.Sin(float64(time.Now().UnixNano()*10^-9) * TimeFactor * eV / hbar)
178			centerY := float32(g.eigenvalues[i] / V0Max)
179			yPrev := float32(g.eigenvectors.At(0, i) * phase)
180			for j := range N - 1 {
181				y := float32(g.eigenvectors.At(j+1, i) * phase)
182
183				vector.StrokeLine(screen,
184					float32(Width)*(float32(j)/float32(N-1)),
185					float32(Height)-(float32(Height)*(centerY)+yPrev),
186					float32(Width)*(float32(j+1)/float32(N-1)),
187					float32(Height)-(float32(Height)*(centerY)+y),
188					2,
189					g.computedColors[i],
190					true)
191				yPrev = y
192			}
193		}
194	}
195
196}
197
198// Layout takes the outside size (e.g., the window size) and returns the (logical) screen size.
199// If you don't have to adjust the screen size with the outside size, just return a fixed size.
200func (g *Game) Layout(outsideWidth, outsideHeight int) (screenWidth, screenHeight int) {
201	return Width, Height
202}
203
204func main() {
205	potential := make([]float64, N)
206	for i := range potential {
207		if i < 3*N/4 && i >= N/4 {
208			potential[i] = 0.1
209		} else {
210			potential[i] = 0.9
211		}
212	}
213	game := &Game{
214		potential:        potential,
215		changedPotential: true,
216	}
217	// Specify the window size as you like. Here, a doubled size is specified.
218	ebiten.SetWindowSize(Width, Height)
219	ebiten.SetWindowTitle("Poisson solver")
220	ebiten.SetTPS(240)
221	// Call ebiten.RunGame to start your game loop.
222	if err := ebiten.RunGame(game); err != nil {
223		log.Fatal(err)
224	}
225
226}