For XLL

solve.py

from scipy.optimize import fsolve
import scipy
import math
import numpy
import sys
# import matplotlib.plot as plt

C = []
S = []

# 生成求解 C1～C4 系列原子的平面坐标的方程组
def composeEqs(circleCenter, A = 1, _lambda = 10):
    keyLength = 0.075186
    def eqs(i):
        x, y = i[0], i[1]
        return [
        A * scipy.sin(2 * math.pi * x / _lambda) - y, 
        (x - circleCenter[0])**2 + (y - circleCenter[1])**2 - keyLength
      ]
    return eqs;

def computePositions(atomCount, A, _lambda):
  loop = atomCount;
  currentLoop = 0;
  atoms = [[0,0]];

  while currentLoop < loop:
    atoms.append(fsolve(composeEqs(atoms[currentLoop], A, _lambda), [atoms[currentLoop][0] + 1, atoms[currentLoop][1]]).tolist())
    currentLoop = currentLoop + 1
  
  return atoms;

def transformPositions(atoms, groupIndex = 1, y = 0.1583):
  result = []
  for atom in atoms:
    result.append([atom[0], y * (groupIndex - 1), atom[1]])
  return result;

def makeGroups(atoms):
  groupA = []
  groupB = []
  for index in range(len(atoms)):
    if index % 2 == 0:
      groupA.append(atoms[index])
    else:
      groupB.append(atoms[index])

  group1 = transformPositions(groupB, groupIndex = 1)
  group2 = transformPositions(groupA, groupIndex = 2)
  group3 = transformPositions(groupB, groupIndex = 3)
  group4 = transformPositions(groupA, groupIndex = 4)

  return [group1, group2, group3, group4]

def getCornorAtoms(sAtomIndexs, cAtomGroups):
  sIndex1 = sAtomIndexs[0]
  sIndex2 = sAtomIndexs[1]
  result = []
  result.append(cAtomGroups[sIndex1][sIndex2])
  result.append(cAtomGroups[sIndex1 - 1][sIndex2 if sIndex1 % 2 != 0 else sIndex2 + 1])
  result.append(cAtomGroups[sIndex1 + 1][sIndex2 if sIndex1 % 2 != 0 else sIndex2 + 1])
  return result;

def composeEqs2(sAtomIndexs, cAtomGroups):
    cornerCAtoms = getCornorAtoms(sAtomIndexs, cAtomGroups)
    keyLength = 0.05827396
    def eqs(i):
        x, y, z = i[0], i[1], i[2]
        return [
        (x - cornerCAtoms[0][0])**2 + (y - cornerCAtoms[0][1])**2 + (z - cornerCAtoms[0][2])**2 - keyLength, 
        (x - cornerCAtoms[1][0])**2 + (y - cornerCAtoms[1][1])**2 + (z - cornerCAtoms[1][2])**2 - keyLength, 
        (x - cornerCAtoms[2][0])**2 + (y - cornerCAtoms[2][1])**2 + (z - cornerCAtoms[2][2])**2 - keyLength, 
      ]
    return eqs;

def computeSAtomPostion(sAtomIndexs, cAtomGroups, A):
  cornerCAtoms = getCornorAtoms(sAtomIndexs, cAtomGroups)
  r1 = fsolve(
    composeEqs2(sAtomIndexs, cAtomGroups), 
    [cornerCAtoms[0][0] + 10 * A, cornerCAtoms[0][1], cornerCAtoms[0][2] + 10 * A]
  ).tolist()
  r2 = fsolve(
    composeEqs2(sAtomIndexs, cAtomGroups), 
    [cornerCAtoms[0][0] - 10 * A, cornerCAtoms[0][1], cornerCAtoms[0][2] - 10 * A]
  ).tolist()
  return [r1, r2]

def computeSAtomPostions(cAtomGroups, A):
  subLoop = len(cAtomGroups[0]);
  result = [[], []];

  for i in [2, 3]:
    for j in range(subLoop):
      result[i - 2].append(computeSAtomPostion([i - 1, j], cAtomGroups, A))

  return result

def doubleCheck(sAtom, cornerCAtoms):
  x = sAtom[0]
  y = sAtom[1]
  z = sAtom[2]
  r1 = (x - cornerCAtoms[0][0])**2 + (y - cornerCAtoms[0][1])**2 + (z - cornerCAtoms[0][2])**2 - 0.05827396
  r2 = (x - cornerCAtoms[1][0])**2 + (y - cornerCAtoms[1][1])**2 + (z - cornerCAtoms[1][2])**2 - 0.05827396
  r3 = (x - cornerCAtoms[2][0])**2 + (y - cornerCAtoms[2][1])**2 + (z - cornerCAtoms[2][2])**2 - 0.05827396
  return [r1, r2, r3]


# atomCount       希望计算出的原子坐标的个数，必填
# A               正弦曲线方程中的系数 A，默认为 1。
# _lambda         正弦曲线方程中的系数 lambda，默认为 10。
# toFile          是否将计算结果或过程写入文本文件，默认为是。
# debug           是否开启 debug 模式，默认为否。开启后，会输出指定索引的 S 原子的坐标的
#                 计算过程和验算结果。
# exampleSIndexs  指定一个 S 原子的索引，用于输出计算过程，仅在 debug 模式下有效，默认
#                 为 [2, 3]。 索引值从 1 开始，以 human-readable 的方式指定。
def runPipeline(atomCount, A = 1, _lambda = 10, toFile = True, debug = False, exampleSIndexs = [2, 3]):
  aFile = 0
  originalOut = 0
  if toFile:
    aFile = open('result.txt', 'w+')
    originalOut = sys.stdout
    sys.stdout = aFile

  cAtomGroups = makeGroups(computePositions(atomCount * 2, A, _lambda))

  if not debug:
    print("四组 C 原子坐标为：")
    for group in cAtomGroups:
      print(group, '\n')
    print('-' * 20, '\n')

    r = computeSAtomPostions(cAtomGroups, A)

    print("两组 S 原子的坐标分别为：")
    for aR in r:
      print(aR, '\n')
  else:
    print("以 S[" + str(exampleSIndexs[0]) + "][" + str(exampleSIndexs[1]) + "] 原子为例，求其坐标。")
    cornerCAtoms = getCornorAtoms(
      [exampleSIndexs[0] - 1, exampleSIndexs[1] - 1], 
      cAtomGroups
    )

    print("相连三个 C 原子坐标为:")
    print(cornerCAtoms)

    print("求解方程，得解：")
    r = fsolve(
      composeEqs2([exampleSIndexs[0] - 1, exampleSIndexs[1] - 1], cAtomGroups), 
      [cornerCAtoms[0][0] + 10 * A, cornerCAtoms[0][1], cornerCAtoms[0][2] + 10 * A]
    ).tolist()
    print("解 1：")
    print(r)
    print("验算：", doubleCheck(r, cornerCAtoms))
    r = fsolve(
      composeEqs2([exampleSIndexs[0] - 1, exampleSIndexs[1] - 1], cAtomGroups), 
      [cornerCAtoms[0][0] - 10 * A, cornerCAtoms[0][1], cornerCAtoms[0][2] - 10 * A]
    ).tolist()
    print("解 2：")
    print(r)
    print("验算：", doubleCheck(r, cornerCAtoms))

  if toFile:
    aFile.close()
    sys.stdout = originalOut


if __name__ == "__main__":
  runPipeline(10, toFile = True, debug = False)
  runPipeline(10, toFile = False, debug = True, exampleSIndexs = [2, 3])