(***********************************************************************
This file was generated automatically by the Mathematica front end.
It contains Initialization cells from a Notebook file, which typically
will have the same name as this file except ending in ".nb" instead of
".m".

This file is intended to be loaded into the Mathematica kernel using
the package loading commands Get or Needs.  Doing so is equivalent to
using the Evaluate Initialiation Cells menu command in the front end.

DO NOT EDIT THIS FILE.  This entire file is regenerated automatically 
each time the parent Notebook file is saved in the Mathematica front end.
Any changes you make to this file will be overwritten.
***********************************************************************)

BeginPackage["SegAnalysis`ZOne`"]

ZOneA::usage =
  "ZOneA[x, xv1, vb1, vh1] gives the expected fraction of A genotype at x, \
when the gene at xv1 affects each viability of zygote.  The relateve \
viabilities of B homozygote and heterozygote to the A homozygote are vb1 and \
vh1, respectively."

ZOneB::usage =
  "VOneB[x, xv1, vb1, vh1] gives the expected fraction of B genotype at x, \
when the gene at xv1 affects each viability of zygote.   The relateve \
viabilities of B homozygote and heterozygote to the A homozygote are vb1 and \
vh1, respectively."

ZOneH::usage =
  "VOneH[x, xv1, vb1, vh1] gives the expected fraction of H genotype at x, \
when the gene at xv1 affects each viability of zygote.   The relateve \
viabilities of B homozygote and heterozygote to the A homozygote are vb1 and \
vh1, respectively."

Begin["`Private`"]

(* === Warning messages ====  *)

SegregAnalysis::lmincrease = 
  "Warning: after `` decreases in the step size, NonlinearFit failed to find \
a new parameter estimate that will decrease the sum of squares."

(* =========    ZOne     =========== *)
(* This uses Guess-Newton method for finding the least-
    squares solution for a single gene segregation distortion model  *)

mnlr[nptsA_, nptsB_, nptsH_, nresponseA_, nresponseB_, nresponseH_, xv1_, 
    vb1_, vh1_, fp_, maxits_, wp_, pg_, ag_, mdls_, zt_] := 
  Module[{predA, predB, predH, chilistA, chilistB, chilistH, tmpchilistA, 
      tmpchilistB, tmpchilistH, tmpchilist, tmpnptsA, tmpnptsB, tmpnptsH, 
      chisq, derAxv1, derBxv1, derHxv1, derxv1, derAvb1, derBvb1, derHvb1, 
      dervb1, derAvh1, derBvh1, derHvh1, dervh1, der, alpha, beta, chisq0, 
      newchi, tmpchi, its, accDelta, NAcc, precDelta, NPrec, invaliddigits, 
      accflag0, accflag1 = False, accflag2 = False, precflag0, 
      precflag1 = False, precflag2 = False, ovflag = False, guessxv1 = xv1,
      guessvb1 = vb1, guessvh1 = vh1, oldxv1, oldvb1, oldvh1, delta, 
      deltaxv1, deltavb1, deltavh1, tmpxv1, tmpvb1, tmpvh1, dits, dmaxits, 
      newpredA, newpredB, newpredH, newchilistA, newchilistB, newchilistH, 
      accgoal = N[10^(-ag)], precgoal = N[10^(-pg)], paramerr, errorDOF},
		
		(* - calculation chi-squars *)
		{{chilistA, chilistB, chilistH}, chisq} = 
      calcchi[{nptsA, nptsB, nptsH}, {nresponseA, nresponseB,  nresponseH}, {
          guessxv1,guessvb1, guessvh1}, wp];
		tmpchi = chisq; newchi = -1; its = 1; 
		oldxv1 = guessxv1;
		oldvb1 = guessvb1; oldvh1 = guessvh1;

		(* main loop *)
		ovflag = Check[
				While[(
						(accDelta = Abs[chisq - newchi];
							NAcc = -Floor[Log[10, accDelta]] - 1;
							precDelta = Abs[chisq - newchi]/(newchi + chisq)/2;
							NPrec = -Floor[Log[10, precDelta]] - 1;
							invaliddigits = 
                  TrueQ[NAcc =!= Infinity && Accuracy[accDelta] < NAcc && 
                      NPrec =!= Infinity && Precision[precDelta] < Nprec];
							!invaliddigits ) &&
							( accflag0 = ( 0 <= accDelta < accgoal) ;
								precflag0 =( 0 <= precDelta < precgoal);
								!((accflag0 && accflag1 && accflag2) || (
                        precflag0 && precflag1 && precflag2))  ) && 
              its <= maxits),
					
					{accflag2, precflag2} = {accflag1, precflag1};
					{accflag1, precflag1} = {accflag0, precflag0};
					chisq = tmpchi;
					oldxv1 = guessxv1; 
					oldvb1 = guessvb1; oldvh1 = guessvh1;
					
					(* === find deltaxv1, deltava1, 
            deltavb1 and deltavh1 by LinearSolve normalmatrix ==== *)
					
					(* === delete not differentiable position ==== *)
					tmpchilistA = Delete[chilistA, Position[nptsA, guessxv1]];
					tmpnptsA =Delete[nptsA, Position[nptsA, guessxv1]];
					tmpchilistB = Delete[chilistB, Position[nptsB, guessxv1]];
					tmpnptsB =Delete[nptsB, Position[nptsB, guessxv1]];
					tmpchilistH = Delete[chilistH, Position[nptsH, guessxv1]];
					tmpnptsH =Delete[nptsH, Position[nptsH, guessxv1]];
					tmpchilist = Join[tmpchilistA , tmpchilistB, tmpchilistH];
					(* calculate jacobian *)
						der={};
					delta={};
					If[!(MemberQ[fp, "xv1"]),
					derAxv1 = N[( Which[
									#1 <= guessxv1, 
									(-8*Sech[2*#1 - 2*guessxv1]^2 + 
                              8*guessvh1*Sech[2*(-#1 + guessxv1)]^2 - 
                              4*(1 + guessvb1 - 2*guessvh1)*
                                Sech[2*(#1 - guessxv1)]^2*
                                Tanh[2*(#1 - guessxv1)])/(
                            4*(1 + guessvb1 + 2*guessvh1)), 
									guessxv1 < #1,
									(-8*(-1 + guessvh1)*Sech[2*(#1 - guessxv1)]^2 - 
                              4*(1 + guessvb1 - 2*guessvh1)*
                                Sech[2*(#1 - guessxv1)]^2*
                                Tanh[2*(#1 - guessxv1)])/(
                            4*(1 + guessvb1 + 2*guessvh1))] &) /@ 
								tmpnptsA, wp];
					derBxv1 = N[( Which[
									#1 <= guessxv1, 
									(-8*guessvb1*Sech[2*#1 - 2*guessxv1]^2 + 
                              8*guessvh1*Sech[2*(-#1 + guessxv1)]^2 - 
                              4*(1 + guessvb1 - 2*guessvh1)*
                                Sech[2*(#1 - guessxv1)]^2*
                                Tanh[2*(#1 - guessxv1)])/(
                            4*(1 + guessvb1 + 2*guessvh1)), 
									guessxv1 < #1,
									(
                            8*(guessvb1 - guessvh1)*
                                Sech[2*(#1 - guessxv1)]^2 - 
                              4*(1 + guessvb1 - 2*guessvh1)*
                                Sech[2*(#1 - guessxv1)]^2*
                                Tanh[2*(#1 - guessxv1)])/(
                            4*(1 + guessvb1 + 2*guessvh1))] &) /@ 
								tmpnptsB, wp];
					derHxv1 = N[( Which[
									#1 <= guessxv1, 
									(-8*guessvh1*Sech[2*#1 - 2*guessxv1]^2 + 
                              4*(1 + guessvb1)*Sech[2*(-#1 + guessxv1)]^2 + 
                              4*(1 + guessvb1 - 2*guessvh1)*
                                Sech[2*(#1 - guessxv1)]^2*
                                Tanh[2*(#1 - guessxv1)])/(
                            2*(1 + guessvb1 + 2*guessvh1)), 
									guessxv1 < #1,
									(-4*(1 + guessvb1 - 2*guessvh1)*Sech[2*(#1 - guessxv1)]^2 + 
                              4*(1 + guessvb1 - 2*guessvh1)*
                                Sech[2*(#1 - guessxv1)]^2*
                                Tanh[2*(#1 - guessxv1)])/(
                            2*(1 + guessvb1 + 2*guessvh1))] &) /@ 
								tmpnptsH, wp];
					derxv1 = Join[derAxv1, derBxv1, derHxv1];
						der=Append[der,derxv1]];
					
					If[!(MemberQ[fp, "vb1"]),
						derAvb1 = N[( Which[
									#1 <= guessxv1, 
									Tanh[2*(#1 - guessxv1)]^2/(4*(1 + guessvb1 + 2*guessvh1)) - (
                              4 + 4*Tanh[2*#1 - 2*guessxv1] + (
                                    1 + guessvb1 - 2*guessvh1)*
                                  Tanh[2*(#1 - guessxv1)]^2 + 
                                4*guessvh1*Tanh[2*(-#1 + guessxv1)])/(
                              4*(1 + guessvb1 + 2*guessvh1)^2), 
									guessxv1 < #1,
									Tanh[2*(#1 - guessxv1)]^2/(4*(1 + guessvb1 + 2*guessvh1)) - (
                              4 + 4*(-1 + guessvh1)*
                                  Tanh[2*(#1 - guessxv1)] + (
                                    1 + guessvb1 - 2*guessvh1)*
                                  Tanh[2*(#1 - guessxv1)]^2)/(
                              4*(1 + guessvb1 + 2*guessvh1)^2)] &) /@ 
								tmpnptsA, wp];
					derBvb1 = N[( Which[
									#1 <= guessxv1, 
									(4 + 4*Tanh[2*#1 - 2*guessxv1] + Tanh[2*(#1 - guessxv1)]^2)/(
                              4*(1 + guessvb1 + 2*guessvh1)) - (
                              4*guessvb1*
                                  Tanh[2*#1 - 2*guessxv1] + (
                                    1 + guessvb1 - 2*guessvh1)*
                                  Tanh[2*(#1 - guessxv1)]^2 + 
                                4*(guessvb1 + 
                                      guessvh1*Tanh[2*(-#1 + guessxv1)]))/(
                              4*(1 + guessvb1 + 2*guessvh1)^2), 
									guessxv1 < #1,
									(4 - 4*Tanh[2*(#1 - guessxv1)] + Tanh[2*(#1 - guessxv1)]^2)/(
                              4*(1 + guessvb1 + 2*guessvh1)) - (
                              4*guessvb1 - 
                                4*(guessvb1 - guessvh1)* 
                                  Tanh[2*(#1 - guessxv1)] + (
                                    1 + guessvb1 - 2*guessvh1)*
                                  Tanh[2*(#1 - guessxv1)]^2)/ (
                              4*(1 + guessvb1 + 2*guessvh1)^2)] &) /@ 
								tmpnptsB, wp];
					derHvb1 = N[( Which[
									#1 <= guessxv1, 
									(-Tanh[2*(#1 - guessxv1)]^2 + 2*Tanh[2*(-#1 + guessxv1)])/(
                              2*(1 + guessvb1 + 2*guessvh1)) - (
                              4*guessvh1*
                                  Tanh[2*#1 - 2*guessxv1] - (
                                    1 + guessvb1 - 2*guessvh1)* 
                                  Tanh[2*(#1 - guessxv1)]^2 + 
                                2*(2*guessvh1 + (1 + guessvb1)*
                                        Tanh[2*(-#1 + guessxv1)]))/(
                              2*(1 + guessvb1 + 2*guessvh1)^2), 
									guessxv1 < #1,
									(2*Tanh[2*(#1 - guessxv1)] - Tanh[2*(#1 - guessxv1)]^2)/(
                              2*(1 + guessvb1 + 2*guessvh1)) - (
                              4*guessvh1 + 
                                2*(1 + guessvb1 - 2*guessvh1)*
                                  Tanh[2*(#1 - guessxv1)] - (
                                    1 + guessvb1 - 2*guessvh1)*
                                  Tanh[2*(#1 - guessxv1)]^2)/(
                              2*(1 + guessvb1 + 2*guessvh1)^2)] &) /@ 
								tmpnptsH, wp];
					dervb1 = Join[derAvb1, derBvb1, derHvb1];
						der=Append[der,dervb1]];
					
					If[!(MemberQ[fp, "vh1"]),
						derAvh1 = N[( Which[
									#1 <= guessxv1, 
									(-2*Tanh[2*(#1 - guessxv1)]^2 + 4*Tanh[2*(-#1 + guessxv1)])/(
                              4*(1 + guessvb1 + 2*guessvh1)) - (
                              4 + 
                                4*Tanh[2*#1 - 2*guessxv1] + (
                                    1 + guessvb1 - 2*guessvh1)*
                                  Tanh[2*(#1 - guessxv1)]^2 + 
                                4*guessvh1*Tanh[2*(-#1 + guessxv1)])/(
                              2*(1 + guessvb1 + 2*guessvh1)^2), 
									guessxv1 < #1,
									(4*Tanh[2*(#1 - guessxv1)] - 2*Tanh[2*(#1 - guessxv1)]^2)/(
                              4*(1 + guessvb1 + 2*guessvh1)) - (
                              4 + 4*(-1 + guessvh1)*
                                  Tanh[2*(#1 - guessxv1)] + (
                                    1 + guessvb1 - 2*guessvh1)*
                                  Tanh[2*(#1 - guessxv1)]^2)/(
                              2*(1 + guessvb1 + 2*guessvh1)^2)] &) /@ 
								tmpnptsA, wp];
					derBvh1 = N[( Which[
									#1 <= guessxv1, 
									(-2*Tanh[2*(#1 - guessxv1)]^2 + 4*Tanh[2*(-#1 + guessxv1)])/(
                              4*(1 + guessvb1 + 2*guessvh1)) - (
                              4*guessvb1*
                                  Tanh[2*#1 - 2*guessxv1] + (
                                    1 + guessvb1 - 2*guessvh1)*
                                  Tanh[2*(#1 - guessxv1)]^2 + 
                                4*(guessvb1 + 
                                      guessvh1*Tanh[2*(-#1 + guessxv1)]))/(
                              2*(1 + guessvb1 + 2*guessvh1)^2), 
									guessxv1 < #1,
									(4*Tanh[2*(#1 - guessxv1)] - 2*Tanh[2*(#1 - guessxv1)]^2)/(
                              4*(1 + guessvb1 + 2*guessvh1)) - (
                              4*guessvb1 - 
                                4*(guessvb1 - guessvh1)*
                                  Tanh[2*(#1 - guessxv1)] + (
                                    1 + guessvb1 - 2*guessvh1)*
                                  Tanh[2*(#1 - guessxv1)]^2)/(
                              2*(1 + guessvb1 + 2*guessvh1)^2)] &) /@ 
								tmpnptsB, wp];
					derHvh1 = N[( Which[
									#1 <= guessxv1, 
									(4 + 4*Tanh[2*#1 - 2*guessxv1] + 2*Tanh[2*(#1 - guessxv1)]^2)/(
                              2*(1 + guessvb1 + 2*guessvh1)) - (
                              4*guessvh1*
                                  Tanh[2*#1 - 2*guessxv1] - (
                                    1 + guessvb1 - 2*guessvh1)*
                                  Tanh[2*(#1 - guessxv1)]^2 + 
                                2*(2*guessvh1 + (1 + guessvb1)*
                                        Tanh[2*(-#1 + guessxv1)]))/(
                                1 + guessvb1 + 2*guessvh1)^2, 
									guessxv1 < #1,
									(4 - 4*Tanh[2*(#1 - guessxv1)] + 2*Tanh[2*(#1 - guessxv1)]^2)/(
                              2*(1 + guessvb1 + 2*guessvh1)) -(
                              4*guessvh1 + 
                                2*(1 + guessvb1 - 2*guessvh1)*
                                  Tanh[2*(#1 - guessxv1)] - (
                                    1 + guessvb1 - 2*guessvh1)*
                                  Tanh[2*(#1 - guessxv1)]^2)/(
                                1 + guessvb1 + 2*guessvh1)^2] &) /@ 
								tmpnptsH, wp];
					dervh1 = Join[derAvh1, derBvh1, derHvh1];
				der=Append[der,dervh1]];
					
					(* normal equation: alpha . deltaguess === beta   *)
					{alpha, beta} = 
            N[{der . Transpose[der], Map[tmpchilist . #1 &, der]}, wp];
					
					 delta = N[LinearSolve[alpha, beta], wp];
					If[!(MemberQ[fp, "xv1"]),deltaxv1=First[delta];delta=Rest[delta], 
            deltaxv1=0];
					If[!(MemberQ[fp, "vb1"]),deltavb1=First[delta]; delta=Rest[delta],  
            deltavb1=0];
					If[!(MemberQ[fp, "vh1"]),deltavh1=First[delta],  deltavh1=0];
					
					{tmpxv1,tmpvb1, tmpvh1} = 
            N[{oldxv1, oldvb1, oldvh1} +  {deltaxv1, deltavb1, deltavh1}, 
              wp];
					
					{{newchilistA, newchilistB, newchilistH}, newchi} = 
            calcchi[{nptsA, nptsB, nptsH}, {nresponseA, nresponseB, 
                nresponseH}, {tmpxv1,tmpvb1, tmpvh1}, wp];
					
					(* find newguess for smaller chisq  *)
					dits = 1; dmaxits = 15;
					While[newchi - chisq >= 0 && dits <= dmaxits,
						{tmpxv1,tmpvb1, tmpvh1} = 
              N[{oldxv1, oldvb1, oldvh1} + 
                  0.5^dits*{deltaxv1, deltavb1, deltavh1}, wp];
						{{newchilistA, newchilistB, newchilistH}, newchi} = 
              calcchi[{nptsA, nptsB, nptsH}, {nresponseA, nresponseB, 
                  nresponseH}, {tmpxv1, tmpvb1, tmpvh1}, wp];
						dits++];
					If[dits > dmaxits, Message[SegregAnalysis::lmincrease, dmaxits]];
					tmpchi = newchi;
					{guessxv1, guessvb1, guessvh1} = {tmpxv1,tmpvb1, tmpvh1};
					chilistA = newchilistA;
					chilistB = newchilistB;
					chilistH = newchilistH;
					printprogress[its, newchi, {guessxv1, guessvb1, guessvh1}];
					its++], (* end of While *)
				True, General::ovfl];  (* end of Check *)
		(* Adjust the accuracy and precision of chisq *)
		If[Accuracy[newchi] > NAcc, newchi = SetAccuracy[newchi, NAcc]];
		If[Precision[newchi] > NPrec, newchi = SetPrecision[newchi, NPrec]];
		errorDOF = 
      Length[chilistA] + Length[chilistB] + Length[chilistH] - 3+Length[fp];
		paramerr = Sqrt[ ( newchi/errorDOF) DiagonalElements[Inverse[alpha]]];
		{newchi, {guessxv1, guessvb1, guessvh1}, paramerr, errorDOF, fp}
	]  (* end ZOne *)

(*   *)
ZOneA[x_, xv1_, vb1_, vh1_] := Which[
		x <= xv1, 
		(4 + 4*Tanh[2*x - 2*xv1] +(1 + vb1 - 2*vh1)*Tanh[2*(x - xv1)]^2 + 
          4*vh1*Tanh[2*(-x + xv1)])/(4*(1 + vb1 + 2*vh1)),
		xv1 < x, 
		(4 + 4*(-1 + vh1)*Tanh[2*(x - xv1)] + (1 + vb1 - 2*vh1)*
            Tanh[2*(x - xv1)]^2)/(4*(1 + vb1 + 2*vh1))]
	
(*  *)
ZOneB[x_, xv1_, vb1_, vh1_] := Which[
		x <= xv1, 
		(4*vb1*Tanh[2*x - 2*xv1] + (1 + vb1 - 2*vh1)*Tanh[2*(x - xv1)]^2 + 
          4*(vb1 + vh1*Tanh[2*(-x + xv1)]))/(4*(1 + vb1 + 2*vh1)),
		xv1 < x, 
		(4*vb1 - 
          4*(vb1 - vh1)*Tanh[2*(x - xv1)] + (1 + vb1 - 2*vh1)*
            Tanh[2*(x - xv1)]^2)/(4*(1 + vb1 + 2*vh1))]
	
(*  *)
ZOneH[x_, xv1_, vb1_, vh1_] := Which[
		x <= xv1, 
		(4*vh1*Tanh[2*x - 2*xv1] -(1 + vb1 - 2*vh1)*Tanh[2*(x - xv1)]^2 + 
          2*(2*vh1 + (1 + vb1)*Tanh[2*(-x + xv1)]))/(2*(1 + vb1 + 2*vh1)), 
		xv1 < x,
		(4*vh1 + 
          2*(1 + vb1 - 2*vh1)*Tanh[2*(x - xv1)] -(1 + vb1 - 2*vh1)*
            Tanh[2*(x - xv1)]^2)/(2*(1 + vb1 + 2*vh1))]

(*   calculation of chi-list and chi-squares *)
calcchi[pts_List, response_List, guess_List, wp_] :=
	Block[{xv1, vb1, vh1, ptsA, ptsB, ptsH, responseA, responseB, responseH, 
      predA, predB, predH, chilistA, chilistB, chilistH, chisq},
		{responseA, responseB, responseH} = response;
		{ptsA, ptsB, ptsH} = pts;
		{xv1, vb1, vh1} =N[guess,wp];
		predA = N[ZOneA[#1, xv1, vb1, vh1] & /@ ptsA, wp];
		predB = N[ZOneB[#1, xv1, vb1, vh1] & /@ ptsB, wp];
		predH = N[ZOneH[#1, xv1, vb1, vh1] & /@ ptsH, wp];
		chilistA = responseA - predA;
		chilistB = responseB - predB;
		chilistH = responseH - predH;
		chisq =N[ 
        Plus @@ (
            Join[Map[#1^2 &, chilistA], Map[#1^2 &, chilistB], 
              Map[#1^2 &, chilistH]]), wp];
		{{chilistA, chilistB, chilistH}, chisq}
		]

(*   *)
printprogress[iteration_, chisq_, params_] := 
     Print[
    StringForm[
      "Iteration:`1` ChiSquared:`2`  {Position, ViabilityB, ViablityH}:`3`", 
											NumberForm[iteration, 6], NumberForm[chisq, 6], 
      NumberForm[params, 6]]]

(*    *)
DiagonalElements[matrix_] := 
  Module[{k, size = Length[matrix]}, 
   Table[matrix[[k,k]], {k, 1, size}]
		]

(* =========================================== *)
End[]
EndPackage[]